Efficient Long-Range Convolutions for Point Clouds

1 INTRODUCTION . Point-cloud representations provide detailed information of objects and environments . The development of novel acquisition techniques , such as laser scanning , digital photogrammetry , light detection and ranging ( LIDAR ) , 3D scanners , structure-from-motion ( SFM ) , among others , has increased the interest of using point cloud representation in various applications such as digital preservation , surveying , autonomous driving ( Chen et al. , 2017 ) , 3D gaming , robotics ( Oh & Watanabe , 2002 ) , and virtual reality ( Park et al. , 2008 ) . In return , this new interest has fueled the development of machine learning frameworks that use point clouds as input . Historically , early methods used a preprocessing stage that extracted meticulously hand-crafted features from the point cloud , which were subsequently fed to a neural network ( Chen et al. , 2003 ; Rusu et al. , 2008 ; Rusu et al. , 2009 ; Aubry et al. , 2011 ) , or they relied on voxelization of the geometry ( Savva et al. , 2016 ; Wu et al. , 2015 ; Riegler et al. , 2017 ; Maturana & Scherer , 2015 ) . The PointNet architecture ( Qi et al. , 2017 ) was the first to handle raw point cloud data directly and learn features on the fly . This work has spawned several related approaches , aiming to attenuate drawbacks from the original methodology , such as PointNet++ ( Qi et al. , 2017 ) , or to increase the accuracy and range of application ( Wang et al. , 2019 ; Zhai et al. , 2020 ; Li et al. , 2018 ; Liu et al. , 2019 ) . Even though such methods have been quite successful for machine learning problems , they rely on an assumption of locality , which may produce large errors when the underlying task at hand exhibits long-range interactions ( LRIs ) . To capture such interactions using standard convolutional layers , one can use wider window sizes , deeper networks , and/or a large number of features , which may increase the computational cost significantly . Several approaches have been proposed to efficiently capture such interactions in tasks such as semantic segmentation , of which the ideas we briefly summarize below . In the multi-scale type of approaches , features are progressively processed and merged . Within this family , there exist several variants , where the underlying neural networks can be either recursive neural networks ( Ye et al. , 2018 ) , convolutional layers ( Xu et al. , 2019 ; Xu et al. , 2018 ) or autoencoders ( Yang et al. , 2018 ; Deng et al. , 2018 ) . Some works have proposed skip connections , following an U-net ( Ronneberger et al. , 2015 ) type architecture ( Zhou & Tuzel , 2018 ; Qi et al. , 2017 ) , while others have focused on using a tree structure for the clustering of the points ( Klokov & Lempitsky , 2017 ; Zeng & Gevers , 2018 ; Gadelha et al. , 2018 ) , or using an reference permutohedral lattices to compute convolutions ( Jampani et al. , 2016 ) whose results are interpolated back to the point cloud ( Su et al. , 2018 ) . Although these methods have been shown to be successful in a range of applications , when the task at hand presents symmetries , such as rotation , translation , and permutation invariance , there is no systematic framework to embed those symmetries into the algorithmic pipelines . Another line of work , relies on interpreting the point cloud as a graph and use spectral convolutions ( Bruna et al . ; Defferrard et al. , 2016 ) , whose cost can scale super-linearly when dealing with LRIs . In applications of machine learning to scientific computing , several classical multilevel matrix factorizations have been rewritten in the context of machine learning ( Kondor et al. , 2014 ) , which have been adapted to handle long-range interactions in the context of end-to-end maps using voxelized geometries in ( Fan et al. , 2019b ; a ; Khoo & Ying , 2019 ; Fan & Ying , 2019 ) resulting in architectures similar to U-nets ( Ronneberger et al. , 2015 ) , which have been extended to point clouds in ( Li et al. , 2020 ) . Due to underlying voxelization of the geometry , it may be difficult for these networks to generalize when the resolution of the voxelization changes . The efficient treatment of LRI for point clouds is also a prominent problem in many physical applications such as molecular modeling and molecular dynamics simulation . While long-range electrostatic interactions are omnipresent , it has been found that effectively short-ranged models can already describe the N -body potential and the associated force field ( Behler & Parrinello , 2007 ; Zhang et al. , 2018a ; b ) for a wide range of physical systems . There have also been a number of recent works aiming at more general systems beyond this regime of effective short-range interactions , such as the work of Ceriotti and co-workers ( Grisafi & Ceriotti , 2019 ; Grisafi et al . ; Nigam et al. , 2020 ; Rossi et al. , 2020 ) , as well as the works of ( Yao et al. , 2018 ; Ko et al. , 2009 ; Hirn et al. , 2017 ; Rupp et al. , 2012 ; Huo & Rupp ; Deng et al. , 2019 ; Bereau et al. , 2018 ; Zhang et al. , 2019 ) . The general strategy is to build parameterized long-range interactions into the kernel methods or neural network models , so that the resulting model can characterize both short-range , as well as long-range electrostatic interactions . In the neural network context , the computational cost of treating the LRIs using these methods can grow superlinearly with the system size . The idea of this work is aligned with the approaches in the molecular modeling community , which constructs a neural network layer to directly describe the LRI . In particular , we present a new longrange convolutional ( LRC ) -layer , which performs a global convolutional operation in nearly-linear time with respect to number of units in the layer . By leveraging the non-uniform Fourier transform ( NUFFT ) ( Dutt & Rokhlin , 1993 ; Greengard & Lee , 2004 ; Barnett et al. , 2019 ) technique , the LRC-layer implements a convolution with a point-wise multiplication in the frequency domain with trainable weights known as Fourier multipliers . The NUFFT is based on the regular fast Fourier transform ( FFT ) ( Cooley & Tukey , 1965 ) with a fast gridding algorithms , to allow for fast convolution on unstructured data . This new LRC-layer provides a new set of descriptors that can seamlessly satisfy relevant symmetries . For instance , when the kernel of the LRI is rotationally invariant , such symmetry can be directly built into the parameterization of the Fourier kernel . Such descriptors can be used in tandem with the descriptors provided by short-range convolutional layers to improve the performance of the neural network . Efficient training of a neural network with the LRC-layer for capturing the information of LRIs is another challenging problem . Short-range models can often be trained with data generated with a relatively small computational box ( called the small-scale data ) , and they can be seamlessly deployed in large-scale systems without significantly increasing the generalization error . On the other hand , long-range models need to be trained directly with data generated in a large computational box ( called the large-scale data ) , and the generation process of such large-scale data can be very expensive . For instance , in molecular modeling , the training data is often generated with highly accurate quantum mechanical methods , of which the cost can scale steeply as O ( Nα ) , where N is the system size and α ≥ 3 . Therefore it is desirable to minimize the number of samples with a large system size . In many applications , the error of the effective short-range model is already modestly small . This motivates us to propose a two-scale training strategy as follows . We first generate many small-scale data ( cheaply and possibly in parallel ) , and train the network without the LRC-layer . Then we use a small number of large-scale data , and perform training with both the short- and long-range convolutional layers . In order to demonstrate the effectiveness of the LRC-layer and the two-scale training procedure , we apply our method to evaluate the energy and force associated with a model N -body potential that exhibit tunable short- and long-range interactions in one , two and three dimensions . The input point cloud consists of the atomic positions , and the output data include the N -body potential , local potential , and the force ( derivative of the N -body potential with respect to atomic positions ) . In particular , the local potential and the force can be viewed as point clouds associated with the atomic positions . The evaluation of the N -body potential is a foundational component in molecular modeling , and LRI plays an important role in the description of ionic systems , macroscopically polarized interfaces , electrode surfaces , and many other problems in nanosciences ( French et al. , 2010 ) . Our result verifies that the computational cost of the long-range layer can be reduced from O ( N2 ) using a direct implementation , to O ( N ) ( up to logarithmic factors ) using NUFFT . Furthermore , we demonstrate that the force , i.e . the derivatives of the potential with respect to all inputs can be evaluated with O ( N ) cost ( up to logarithmic factors ) . In terms of sample efficiency , we find that for the model problem under study here , the two-scale training strategy can effectively reduce the number of large-scale samples by over an order of magnitude to reach the target accuracy . This can be particularly valuable in the context of molecular modeling , where accurate data are often obtained from first principle electronic structure calculations . Such calculations are often very expensive for large scale systems , and the number of large-scale samples is thus limited . 2 LONG-RANGE CONVOLUTIONAL LAYER . Convolutional layers are perhaps the most important building-block in machine learning , due to their great success in image processing and computer vision . A convolutional layer convolves the input , usually an array , with a rectangular mask containing the trainable parameters . When the mask can be kept small ( for example while extracting localized features ) , the convolution layer is highly efficient and effective . A different way for computing a convolution is to use the convolutional theorem as follows : ( 1 ) compute the Fourier transform of the input , ( 2 ) multiply with the Fourier transform of the mask , i.e.m the Fourier multiplier , and ( 3 ) inverse Fourier transform back . In this case , the trainable parameters are the DOFs of the Fourier multipliers and the Fourier transforms are computed using the fast Fourier transform ( FFT ) . This alternative approach is particularly attractive for smooth kernels with large support ( i.e. , smooth long-range interactions ) because the computational cost does not increase with the size of the mask . To the best of our knowledge , this direction has not been explored for LRIs and below we detail now to apply this to point clouds . Given a point cloud { xi } Ni=1 ⊂ Rd and scalar weights { fi } Ni=1 , we consider the problem of computing the quantity ui : = ∑N j=1 φθ ( xi − xj ) fj at each xi . Here the function φθ ( · ) is the kernel with a generic trainable parameter θ . At first glance the cost of this operation scales as O ( N2 ) : we need to evaluate ui for each point xi , which requiresO ( N ) work per evaluation . By introducing a generalized function f ( y ) = ∑ i fi · δ ( y − xi ) and defining a function u ( x ) = ∫ φθ ( x− y ) f ( y ) dy , one notices that ui is the value of u ( x ) at x = xi . The advantage of this viewpoint is that one can now invoke the connection between convolution and Fourier transform û ( k ) = φ̂θ ( k ) · f̂ ( k ) , ( 1 ) where φ̂θ ( k ) is a trainable Fourier multiplier . This approach is suitable for point clouds since the trainable parameters are decoupled from the geometry of the point cloud . To make this approach practical , one needs to address two issues : ( 1 ) the non-uniform distribution of the point cloud and ( 2 ) how to represent the multiplier φ̂θ ( k ) . Non-uniform distribution of the point cloud Equation 1 suggests that one can compute the convolution directly using the convolution theorem , which typically relies on the FFT to obtain a low-complexity algorithm . Unfortunately , { xi } Ni=1 do not form a regular grid , thus FFT can not be directly used . We overcome this difficulty by invoking the NUFFT1 ( Dutt & Rokhlin , 1993 ) , which serves as the corner-stone of our instance of the LRC-layer2 . 1See Appendix C.2 for further details . 2We point out , that one could in practice use an fast summation algorithm , such as the fast multipole method ( FMM ) introduced by Greengard & Rokhlin ( 1987 ) , to evaluate ui . This would results in the same complexity if Algorithm 1 Long-range convolutional layer Input : { xi } Ni=1 , { fi } Ni=1 Output : { xi } Ni=1 , { ui } Ni=1 , where ui = ∑N j=1 fjφθ ( xi − xj ) . 1 : Define the generalized function : f ( x ) = ∑N j=1 fjδ ( x− xj ) 2 : Mollify the Dirac deltas : fτ ( x ) = ∑N j=1 fjgτ ( x− xj ) , where gτ is defined in Appendix C.2 3 : Sample in a regular grid : fτ ( x ` ) = ∑N j=1 gτ ( x ` − xj ) for x ` in grid of size LFFT in each dim 4 : Compute FFT : Fτ ( k ) = FFT ( fτ ) ( k ) 5 : Re-scale the signal : F ( k ) = √ π τ e k2τFτ ( k ) 6 : Multiply by Fourier multipliers : v̂ ( k ) = φ̂θ ( k ) · F ( k ) 7 : Re-scale the signal : v̂−τ ( k ) = √ π τ e k2τ v̂ ( k ) 8 : Compute IFFT : u−τ ( x ` ) = IFFT ( v̂−τ ) ( x ) for x ` on the regular grid 9 : Interpolate to the point cloud : ui = u ( xi ) = u−τ ∗ gτ ( xi ) The LRC-layer is summarized in Alg . 1 , where τ is chosen following Dutt & Rokhlin ( 1993 ) . The inputs of this layer are the point cloud { xi } Ni=1 and the corresponding weights { fi } Ni=1 . The outputs are ui ≡ u ( xi ) for i = 1 , ... , N . The number of elements in the underlying grid NFFT = LdFFT is chosen such that the kernel is adequately sampled and the complexity remains low . As shown in Appendix C.5 , one only needs a relatively small LFFT . Even though the precise number is problemspecific , given that the goal is to approximate LRIs that are supposedly smooth , it can be captured with a relatively small number of Fourier modes . The LRC-layer is composed of three steps : ( 1 ) It computes the Fourier transform from the point cloud to a regular grid using the NUFFT algorithm ( lines 2− 5 in Alg . 1 and showcased in Fig . 2 ) . ( 2 ) It multiplies the result by a set of trainable Fourier multipliers ( line 6 in Alg . 1 ) . ( 3 ) It computes the inverse Fourier transform from the regular grid back to the point cloud ( lines 7− 9 in Alg . 1 ) . Within the LRC-layer in Alg . 1 , the only trainable component is the parameter θ of the Fourier multiplier φ̂θ ( k ) . The remaining components , including the mollifier gτ ( · ) and the Cartesian grid size , are taken to be fixed . One can , in principle , train them as well , but it comes with a much higher cost . Among the steps of Alg . 1 , the sampling operator , the rescaling operator , the interpolation operator , and the Fourier transforms , are all linear and non-trainable . Therefore , derivative computations of backpropagation just go through them directly . Alg . 1 is presented in terms of only one single channel or feature dimension , i.e. , fj ∈ R and ui ∈ R. However , it can be easily generalized to multiple channels , for example fj ∈ Rd1 and ui ∈ Rd2 . In this case , the Fourier multiplier φ̂θ ( k ) at each point k is a d2 × d1 matrix , and all Fourier transforms are applied component-wise . Representation of the Fourier multiplier A useful feature of the LRC-layer is that it is quite easy to impose symmetries on the Fourier multipliers . For example , if the convolution kernel φθ ( · ) is constrained to have parity symmetry , rotational symmetry , smoothness or decay properties , these constraints can be imposed accordingly on the coefficients of the Fourier multipliers φ̂θ ( k ) . When the size of the training data is limited , it is often necessary to reduce the number of trainable parameters in order to regularize the kernel . For example , we may parameterize the Fourier multiplier as a linear combination of several predetermined functions on the Fourier grid . This is the procedure used in molecular modeling ( Grisafi & Ceriotti , 2019 ; Yao et al. , 2018 ; Ko et al. , 2009 ) , and also in our numerical examples in equation 7 . We also remark that the LRC-layer described here can be applied to point clouds a way similar to a standard convolution layer applied to images and multiple LRC-layers can be composed on top of each other . the kernel is fixed . However , in order for the kernel to be trainable , this would require a different algorithm for each iteration , including the computation of the derivatives , thus increasing the computational cost and rendering the implementation significantly more cumbersome . Under review as a conference paper at ICLR 2021 < latexit sha1_base64= '' CzgOsCb5dn9rKCMHR0qbPmNYaK8= '' > AAAB6HicbVBNS8NAEJ3Ur1q/qh69LBbBU0lEsceCF48t2A9oQ9lsJ+3azSbsboQS+gu8eFDEqz/Jm//GbZuDtj4YeLw3w8y8IBFcG9f9dgobm1vbO8Xd0t7+weFR+fikreNUMWyxWMSqG1CNgktsGW4EdhOFNAoEdoLJ3dzvPKHSPJYPZpqgH9GR5CFn1FipGQ7KFbfqLkDWiZeTCuRoDMpf/WHM0gilYYJq3fPcxPgZVYYzgbNSP9WYUDahI+xZKmmE2s8Wh87IhVWGJIyVLWnIQv09kdFI62kU2M6ImrFe9ebif14vNWHNz7hMUoOSLReFqSAmJvOvyZArZEZMLaFMcXsrYWOqKDM2m5INwVt9eZ20r6redfWmeV2p1/I4inAG53AJHtxCHe6hAS1ggPAMr/DmPDovzrvzsWwtOPnMKfyB8/kDyoeM6A== < /latexit > < latexit sha1_base64= '' cf2V4AnIhRim/3JeHlwyQomsm9U= '' > AAAB6HicbVDLTgJBEOzFF+IL9ehlIjHxRHYNRo4kXjxCIo8ENmR26IWR2dnNzKyREL7AiweN8eonefNvHGAPClbSSaWqO91dQSK4Nq777eQ2Nre2d/K7hb39g8Oj4vFJS8epYthksYhVJ6AaBZfYNNwI7CQKaRQIbAfj27nffkSleSzvzSRBP6JDyUPOqLFS46lfLLlldwGyTryMlCBDvV/86g1ilkYoDRNU667nJsafUmU4Ezgr9FKNCWVjOsSupZJGqP3p4tAZubDKgISxsiUNWai/J6Y00noSBbYzomakV725+J/XTU1Y9adcJqlByZaLwlQQE5P512TAFTIjJpZQpri9lbARVZQZm03BhuCtvrxOWldlr1K+blRKtWoWRx7O4BwuwYMbqMEd1KEJDBCe4RXenAfnxXl3PpatOSebOYU/cD5/AOXPjPo= < /latexit > < latexit sha1_base64= '' 9sT+zsi99ZHLmWbkwsJA3sBXKbw= '' > AAACHHicbVDLSgMxFM34rPU16tJNsAgVpMxoxW6EghtXUsE+oFNLJs20oZkHyR2xDPMhbvwVNy4UceNC8G/MtF1o64VwDufcy809biS4Asv6NhYWl5ZXVnNr+fWNza1tc2e3ocJYUlanoQhlyyWKCR6wOnAQrBVJRnxXsKY7vMz85j2TiofBLYwi1vFJP+AepwS01DVPHcE8cJIxFL1uwtNj/JCBI3l/AEcTcFItXdjpXXKd5rtmwSpZ48LzxJ6SAppWrWt+Or2Qxj4LgAqiVNu2IugkRAKngqV5J1YsInRI+qytaUB8pjrJ+LgUH2qlh71Q6hcAHqu/JxLiKzXyXd3pExioWS8T//PaMXiVTsKDKAYW0MkiLxYYQpwlhXtcMgpipAmhkuu/YjogklDQeWYh2LMnz5PGSckul85uyoVqZRpHDu2jA1RENjpHVXSFaqiOKHpEz+gVvRlPxovxbnxMWheM6cwe+lPG1w9426Iw < /latexit > < latexit sha1_base64= '' OigvCa5o4zDt9IMj06BI+xVCTFQ= '' > AAACGnicbVDLSgMxFM3UV62vUZdugkVoF5YZqdhNoeDGlVSwD2jHIZNm2rSZB0lGWob5Djf+ihsXirgTN/6NaTuIth4IHM45l5t7nJBRIQ3jS8usrK6tb2Q3c1vbO7t7+v5BUwQRx6SBAxbwtoMEYdQnDUklI+2QE+Q5jLSc0eXUb90TLmjg38pJSCwP9X3qUoykkmzddO24K1GUFMZFWIVdEXl2PKyayV18nUDXHsL+T+AUju1h0dbzRsmYAS4TMyV5kKJu6x/dXoAjj/gSMyRExzRCacWIS4oZSXLdSJAQ4RHqk46iPvKIsOLZaQk8UUoPugFXz5dwpv6eiJEnxMRzVNJDciAWvan4n9eJpFuxYuqHkSQ+ni9yIwZlAKc9wR7lBEs2UQRhTtVfIR4gjrBUbeZUCebiycukeVYyy6Xzm3K+VknryIIjcAwKwAQXoAauQB00AAYP4Am8gFftUXvW3rT3eTSjpTOH4A+0z2+P5Z/q < /latexit > < latexit sha1_base64= '' xBA9/PTg3eHQ5avpSkVxfJfaLOc= '' > AAACJXicbVDLSgMxFM3UV62vUZdugkVoF5YZqdiFhYIbV1LBPqAzDpk006bNPEgy0jLMz7jxV9y4sIjgyl8xfSy09UDgcM653NzjRowKaRhfWmZtfWNzK7ud29nd2z/QD4+aIow5Jg0cspC3XSQIowFpSCoZaUecIN9lpOUOb6Z+64lwQcPgQY4jYvuoF1CPYiSV5OjXnpNYEsVpYaQIYSwtwiq0ROw7yaBqpo/JXQp7yxl4DkfOoOjoeaNkzABXibkgebBA3dEnVjfEsU8CiRkSomMakbQTxCXFjKQ5KxYkQniIeqSjaIB8IuxkdmUKz5TShV7I1QsknKm/JxLkCzH2XZX0keyLZW8q/ud1YulV7IQGUSxJgOeLvJhBGcJpZbBLOcGSjRVBmFP1V4j7iCMsVbE5VYK5fPIqaV6UzHLp8r6cr1UWdWTBCTgFBWCCK1ADt6AOGgCDZ/AK3sFEe9HetA/tcx7NaIuZY/AH2vcPOzWlCQ== < /latexit > < latexit sha1_base64= '' 0Pdhr03PRt6u7uFQmsGCy911yOY= '' > AAACFXicbVDLSgMxFM3UV62vqks3wSK0IGVGKnYjFITiskJf0JYhk2ba0MyD5I5YhvkJN/6KGxeKuBXc+Tdm2i609UDC4ZxzSe5xQsEVmOa3kVlb39jcym7ndnb39g/yh0dtFUSSshYNRCC7DlFMcJ+1gINg3VAy4jmCdZzJTep37plUPPCbMA3ZwCMjn7ucEtCSnT+v230gUXFSwte4D+wB4tkNENfrzSQpunacBpKSjtj5glk2Z8CrxFqQAlqgYee/+sOARh7zgQqiVM8yQxjERAKngiW5fqRYSOiEjFhPU594TA3i2VYJPtPKELuB1McHPFN/T8TEU2rqOTrpERirZS8V//N6EbjVQcz9MALm0/lDbiQwBDitCA+5ZBTEVBNCJdd/xXRMJKGgi8zpEqzllVdJ+6JsVcqXd5VCrbqoI4tO0CkqIgtdoRq6RQ3UQhQ9omf0it6MJ+PFeDc+5tGMsZg5Rn9gfP4AXkueTw== < /latexit > < latexit sha1_base64= '' s7++NogcTuVwzYAFvVpPjIhc0y4= '' > AAACIHicbVDLSgMxFM34rPVVdekmWIS6KTNF0Y1QFIrLCrYVOrVk0kwbmnmY3BFKmE9x46+4caGI7vRrTNtZaOuBwMk595Dc48WCK7DtL2thcWl5ZTW3ll/f2NzaLuzsNlWUSMoaNBKRvPWIYoKHrAEcBLuNJSOBJ1jLG16O/dYDk4pH4Q2MYtYJSD/kPqcEjNQtnNZKw6NzV91L0K4vCdVuzFPtAknSFLM7PbzTlRRP7rjWnRom0i0U7bI9AZ4nTkaKKEO9W/h0exFNAhYCFUSptmPH0NFEAqeCpXk3USwmdEj6rG1oSAKmOnqyYIoPjdLDfiTNCQFP1N8JTQKlRoFnJgMCAzXrjcX/vHYC/llH8zBOgIV0+pCfCAwRHreFe1wyCmJkCKGSm79iOiCmJjCd5k0JzuzK86RZKTvH5ZPr42L1Iqsjh/bRASohB52iKrpCddRAFD2iZ/SK3qwn68V6tz6mowtWltlDf2B9/wBJ3KOs < /latexit >

Summary: The paper proposes an efficient long-range convolution method for point clouds by using the non-uniform Fourier transform. The long-range convolutional (LRC)-layer mollifies the point cloud to an adequately sized regular grid, computes its Fourier transform, multiplies the results by a set of trainable Fourier multipliers, computes the inverse Fourier transform, and finally interpolates the result back to the point cloud. The method is demonstrated to be effective by a N-body problem.