LSH-SMILE: Locality Sensitive Hashing Accelerated Simulation and Learning

The advancement of deep neural networks over the last decade has enabled progress in scientific knowledge discovery in the form of learning Partial Differential Equations ( PDEs ) directly from experiment data . Nevertheless , forward simulation and backward learning of large-scale dynamic systems requires handling billions mutually interacting elements , the scale of which overwhelms current computing architectures . We propose Locality Sensitive Hashing Accelerated Simulation and Learning ( LSH-SMILE ) , a unified framework to scale up both forward simulation and backward learning of physics systems . LSH-SMILE takes advantages of ( i ) the locality of PDE updates , ( ii ) similar temporal dynamics shared by multiple elements . LSH-SMILE hashes elements with similar dynamics into a single hash bucket and handles their updates at once . This allows LSH-SMILE to scale with respect to the number of non-empty hash buckets , a drastic improvement over conventional approaches . Theoretically , we prove a novel bound on the errors introduced by LSH-SMILE . Experimentally , we demonstrate that LSH-SMILE simulates physics systems at comparable quality with exact approaches , but with way less time and space complexity . Such savings also translate to better learning performance due to LSH-SMILE ’ s ability to propagate gradients over a long duration . 1 Introduction . Learning-driven scientific discovery has enjoyed rapid progress thanks to the advancement of deep neural networks over the last decade . Since Partial Differential Equations ( PDEs ) are widely used to model physics systems , a fruitful line of research has been developed focusing on learning PDEs from experimental data , including Finzi et al . [ 2020 ] , Greydanus et al . [ 2019 ] , Matsubara et al . [ 2020 ] . Nevertheless , successful applications of machine learning for scientific discovery still face multiple challenges , many of which are computational . Both the forward simulation and backward learning of large-scale dynamic systems requires handling billions mutually interacting elements , the scale of which overwhelms current computing architectures . Forward simulation involves simulating the trajectory of a PDE system from a given starting state . Recent work of Greydanus et al . [ 2019 ] demonstrates such process can be represented using a neural network similar to ResNet . Backward learning is to discover ( the parameters of ) the physics model automatically from experimental data , which also attracts recent attention in Niu et al . [ 2020 ] and Xue et al . [ 2021 ] . Backward learning can be achieved by embedding a neural network modeling the forward simulation into the overall architecture and minimizing a loss function which penalizes the difference between the simulated result and the observed data via back-propagation . In both forward simulation and backward learning , billions of mutually interacting elements resulted from applying the finite difference or finite element 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . approaches have to be handled efficiently . Current brute-force approaches which treat each element as one matrix element do not scale to meet the computational need . To tackle the computational bottleneck , we propose Locality Sensitive Hashing Accelerated Simulation and Learning ( LSH-SMILE ) , a unified framework to scale up both forward simulation and backward learning of physics systems represented in a set of PDEs . The method is based on three key observations . The first is that only a small fraction of elements change while the majority of elements remain the same in one step forward simulation . The second is that the one-step update of one element depends on the values of a small set of neighboring elements . The third is that many elements share similar temporal dynamics . The three observations are shared among the simulations of a wide variety of physics phenomena , especially in the so-called interface problems that have applications in fluid dynamics , heat transfer , cracking formation , etc . Motivated by these observations , LSH-SMILE harnesses Locality Sensitive Hashing ( LSH ) to boost the computational throughput . By representing each element as a vector of its value and its neighbor elements ’ value , we hash elements into hash buckets via LSH , where elements with similar self and neighboring element values are grouped into the same bucket . Then the forward simulation and backward learning are carried out for all the elements hashed into the same bucket at once . In this way , the time and space complexities of LSH-SMILE are reduced to be propotional to the number of non-empty hash buckets , a drastic improvement compared to the number of distinct elements in the brute-force algorithm . We also prove a novel bound on the quality of the approximation of LSH-SMILE and provide an interface to control the error rate . LSH-SMILE explores a novel way of using LSH in learning and simulation . Compared to traditional LSH usage case in the nearest neighbor search where the hash table stays the same , LSH-SMILE needs to maintain a hash table that changes dynamically over time . New theoretical and experimental ideas can be sparkled via exploring this new usage of LSH . Experimentally , LSH-SMILE is able to simulate and learn physics systems represented in PDEs at a comparable quality with the exact methods while saving drastically in computation . We focus our attention on systems in nano-physics , where we simulate and learn the grain growth in materials ( Fan and Chen [ 1997 ] ) and model the spatial temporal dynamics of void shaped defects , namely nanovoids , in materials under high temperature and irradiation ( Millett et al . [ 2011 ] ) . The proposed LSH-SMILE reduces the overall simulation time compared to an exact Torch implementation by 70 % in the grain growth simulation , and by 95 % in the nanovoid simulation . In terms of precision , LSH-SMILE matches the ground truth results to the scale of 10−5 in both the nanovoid and the grain growth simulations after 100 simulation steps . Such savings directly translate to better learning performance . In our experiment which uses LSH-SMILE to learn the physics parameters governing a grain growth process , LSH-SMILE-based learning approach can successfully identify the correct model parameters while brute-force baseline approaches can not , The success is due to a longer forward simulation ( 30 time steps ) embedded in the backward learning process enabled by the computational savings of LSH-SMILE . Baseline approaches can only embed 10 forward simulation steps with similar computational budget . One related approach is the Fast Multi-pole Method ( Rokhlin [ 1985 ] ) used in astrophysics . FMM is used to calculate long-ranged forces in the n-body problem , by grouping elements which are close in distance into a single source . However , LSH-SMILE groups elements based on their similarity in future updates . We believe this is a key difference . Overall , the computational innovation of LSH-SMILE opens the door for faster and better scientific discoveries . 2 Preliminaries . Energy-based systems represented in Partial Differential Equations ( PDEs ) . The dynamics of many physics systems can be described with first or second-order Partial Differential Equations ( PDEs ) modeling the system energy . Given the state ~u ( ~p , t ) = ( u1 ( ~p , t ) , . . . , um ( ~p , t ) ) T , where ~p = ( x1 , . . . , xd ) T denotes the spatial coordinates and t denotes time , such PDEs are in the form : ∂~u ( ~p , t ) ∂t = D ( ~u ) +G ( ~u ) ∇F ( ~u ) + I ( ~u ) ∇2H ( ~u ) . ( 1 ) Here , F and H are functions that map Rm to R. ∇ = ( ∂∂x1 , . . . , ∂ ∂xd ) T represents the gradient operator , ∇2 = ∇ · ∇ is the Laplace operator . D is a function that maps Rm to Rm . G ( ~u ) is an m-by-d matrix and I ( ~u ) is a m-by-1 vector . Many notable physics systems can be represented using PDEs in Equation 1 . The widely used AllenCahn and Cahn-Hilliard Equations in phase field modeling are good examples . The Cahn–Hilliard equation and Allen-Cahn equation are used to describes the physics process of field variables in nano-scale materials . The values of the field variables represent the micro-structure composition at different spatial coordinates . The Cahn–Hilliard equation has the general form : ∂u ∂t = ∇ · ( M∇ 1 N δF δu ) . ( 2 ) The Allen-Cahn equation has the general form : ∂v ∂t = −L δF δv . ( 3 ) Here u and v are the field variables of interest and they are assumed to be continuous and changing rapidly across the inter-facial regions . t is time . M , N , L are physics parameters that are related to the changing process . F is the energy function which is different in different applications . More details are left in the supplementary materials . Discretization with finite difference . Finite difference methods allow us to simulate and learn PDEs using the forward and backward propagation of a neural network . It has attracted recent attentions in recent works , e.g. , Wang et al . [ 2020 ] , Dong and Simos [ 2017 ] , Forsythe and Wasows [ 1963 ] and Shojaei et al . [ 2019 ] . More specifically , let { xi,1 , xi,2 , . . . , xi , R } be a finite discretization of the domain of xi , where each xi , k is called a discretization point . { t1 , t2 , . . . , tS } is the discretization of the time . Let ~i = ( i1 , i2 , . . . , id ) . We use ~u ( ~i , j ) as a shorthand for ~u evaluated at position ~i = ( x1 , i1 , . . . , xd , id ) T and time tj , i.e. , ~u ( ~i , j ) = ~u ( x1 , i1 , . . . , xd , id , tj ) . Each discretized position ~i is also called an element . The left-hand side of Equation 1 can be discretized using the finite difference : ∂~u ( ~i , j ) ∂t ≈ ~u ( ~i , j+1 ) −~u ( ~i , j ) tj+1−tj . Similarly , notice ∇F ( ~u ) = ( ∂F ( ~u ) ∂x1 , . . . , ∂F ( ~u ) ∂xd ) T , where ∂F ( ~u ) ∂xl = ∑m k=1 ∂F ( ~u ) ∂uk ∂uk ∂xl . We need ∂uk ( ~i , j ) ∂xl to compute ∇F and ∂uk ( ~i , j ) ∂xl can be approximated by uk ( ~i+1l , j ) −uk ( ~i , j ) xl , il+1−xl , il . Here , ~i + 1l means to move the l-th coordinate of~i , namely il , to the next discretization point il+1 . Repeating this type of calculations for the second-order derivatives∇2H ( ~u ) , the PDE in Equation 1 can be approximated in the following general form : ~u ( ~i , j + 1 ) = ~u ( ~i , j ) + δtQ ( { ~u ( ~i′ , j ) , ~i′ ∈ N ( ~i ) } ) . ( 4 ) Here , we intentionally use function Q to abstract out the actual form , because its derivation is pure arithmetic and is only marginally related to the main purpose of this paper . One important fact to notice : Q depends on a small set of ~u ( ~i′ , j ) ’ s , in which~i′ is a neighboring element of~i . We use N ( ~i ) to represent the set of neighbors of~i . In the two applications we consider in this paper , N ( ~i ) is actually small and only consists of the first and second order neighbors . Two elements are first or second order neighbors if and only if they differ in one dimension within two discretization points , or they differ in two dimensions , each dimension within one discretization point . Forward simulation . By discretizing PDEs with finite difference approach ( Equation 4 ) , we can simulate future states of a physics system from an initial state by updating the left-hand side with values from the right-hand side of Equation 4 repeatedly . Interestingly , such a process can be implemented as a multi-layer convolutional neural network ( see , e.g. , Xue et al . [ 2021 ] ) . See Figure 1 for a demonstration of the forward simulation process . Here we use the grain growth as an example . Backward Learning . Backward learning allows us to learn the parameters of the PDE that governs the dynamics of a physics system from experiment data . It is developed with a series of work , e.g . Wen et al . [ 2021 ] , Dupont et al . [ 2019 ] , Xue et al . [ 2021 ] and Rubanova et al . [ 2019 ] . During backward learning , physics rules are expressed in the form of Equation 4 initialized with random parameters . The dataset consists of the observed pairs of states : ~u ( ~p , t ) and ~u ( ~p , t + T ) , where ~u ( ~p , t ) is the state of physics system at time t and ~u ( ~p , t+ T ) is the state at time t+ T . Here T is a constant set manually . Starting with ~u ( ~p , t ) , a repeated evaluation of Equation 4 of T times will yield the simulated state ~u′ ( ~p , t + T ) . A loss function is defined to penalize the difference between the simulated state ~u′ ( ~p , t+ T ) and observed state ~u ( ~p , t+ T ) . In our experiment , we use the L2 loss function . Back-propagation ( Rumelhart et al . [ 1986 ] , Robbins and Monro [ 1951 ] ) is then applied to minimize the difference . Upon convergence , correct physics parameters are learned which yield the same temporal dynamics as the observed data . Locality Sensitive Hashing ( LSH ) . For a domain S with distance measure D , a LSH family is : Definition 1 . H = { h : S → U } is called a ( r1 , r2 , p , q ) -sensitive LSH function family for D if for any two points x , y ∈ S , one function h chosen uniformly at random fromH satisfies : • if D ( x , y ) ≤ r1 , then Ph∈H [ h ( x ) =h ( y ) ] ≥ p , • if D ( x , y ) ≥ r2 , then Ph∈H [ h ( x ) =h ( y ) ] ≤ q . In this paper , Euclidean distance is used as the distance measure D. In this case , a function in the LSH hash function family has the following form ha , b ( v ) = ⌊ a·v+b r ⌋ , where a is a d-dimensional vector with entries chosen independently at random from a standard Gaussian distribution and b is a real number chosen uniformly from the range [ 0 , r ] . r is a hyper-parameter denoting the size of the hash bucket . We refer the reader to Datar et al . [ 2004 ] for the selection of r and the resulting performance ( i.e. , the values of r1 , r2 , p , q ) .

Summary: This paper proposes a novel LSH based approach for efficient forward and backward pass of Partial Differential Equations (PDEs). The proposed LSH-SMILE approach focus on local updates and use LSH to find the local neighbors. Moreover, LSH-SMILE propose a bucket merging method to avoid empty buckets.