Generalized Fourier Features for Coordinate-Based Learning of Functions on Manifolds

1 INTRODUCTION . Recent breakthroughs on learning representations of 3D shapes ( Mescheder et al. , 2019 ; Park et al. , 2019 ; Sitzmann et al. , 2019 ) or scenes ( Mildenhall et al. , 2020 ) employ the so-called “ coordinatebased ” networks , which take low-dimensional coordinates as inputs and approximate a continuous function . These are sometimes called “ implicit models ” , when the approximated function implicitly represents the desired output ; a typical example is using a signed distance function to represent a 3D shape ( Park et al. , 2019 ) . Perhaps the most important recent advancement in this line of research is NeRF ( Mildenhall et al. , 2020 ) . One of the reasons for NeRF ’ s impressive performance is the positional encoding of input coordinates using sinusoidals of various frequencies , a technique that has been widely adopted ( Liu et al. , 2020 ; Schwarz et al. , 2020 ; Yariv et al. , 2020 ) and studied ( Tancik et al. , 2020 ; Zheng et al. , 2021 ) . The sinusoidals typically used for positional encoding are elements of orthonormal bases for functions on Euclidean spaces . Our key observation is that , to generalize this idea to non-Euclidean manifolds , we should use orthonormal basis functions on the manifold . Figure 1 ( middle ) illustrates the evaluation of the Euclidean basis functions on spherical coordinates , which break orthogonality , results in uneven frequency distribution and singularities near the poles . The spherical harmonics are orthonormal basis functions for the sphere and do not exhibit these undesired properties . NeRF ( Mildenhall et al. , 2020 ) and subsequent works avoid the problem just described by parametrizing the view direction as a unit vector in R3 instead of two angles . We argue that this overparametrization unnecessarily increases the dimensionality of the problem , since the number of basis functions for a fixed bandwidth b grows exponentially with the number of dimensions d. Another relevant example is IPDF ( Murphy et al. , 2021 ) , which takes inputs on the 3-dimensional manifold SO ( 3 ) in the form of a flattened rotation matrix ( a vector in R9 ) , and applies positional encoding by sparsely sampling the space of O ( b9 ) basis elements . We show that our approach improve results for both NeRF and IPDF . Our basic idea is to find an orthonormal basis for the space of functions on the manifold of interest , and introduce principled methods for choosing a subset of the ( possibly infinite ) basis elements on which the input coordinates are evaluated . We describe mathematical techniques to obtain bases for large classes of manifolds and show experiments on a variety of them . Our main contributions are : • We introduce a principled way to apply positional encoding for coordinate-based learning of high-frequency functions on manifolds . • We prove that our approach is shift-invariant under the light of the neural tangent kernel ( NTK ) theory ( Jacot et al. , 2018 ; Tancik et al. , 2020 ) and for the appropriate “ shift ” on the manifold , which implies a manifold convolutional behavior . • Our experiments show the advantages of the proposed methodology on different applications and manifolds : 1 ) learning panoramas on the sphere , 2 ) learning probability distributions on the rotation manifold , 3 ) learning neural radiance fields on the product of cube and sphere , and 4 ) learning light fields represented as the product of spheres . 2 RELATED WORK . The idea of positional encoding using a Fourier basis appeared as early as Rahimi & Recht ( 2008 ) . They introduced the so-called random Fourier features to accelerate training of kernel methods . The idea is to approximate a shift-invariant kernel with random directions sampled from its Fourier transform , on which the inputs are evaluated . More recently , positional encoding has become popular for sequence modeling in natural language processing and often appears in attention layers and transformers as a way to encode the order of the input tokens . Gehring et al . ( 2017 ) use a learned embedding while Vaswani et al . ( 2017 ) use the Fourier basis for encoding input token positions . Alternative approaches were introduced by Xu et al . ( 2019 ) ; Wang et al . ( 2020 ) . Rahaman et al . ( 2019 ) demonstrated that neural networks tend to learn low frequencies more easily than high frequencies , a property called “ spectral bias ” . It also shows , for a simple task of approximating an 1D function , that projecting the input coordinate into a basis of sinusoids facilitates learning high frequencies . NeRF ( Mildenhall et al. , 2020 ) reproduced this finding in the challenging task of photorealistic novel view synthesis from a collection of images . They showed that positional encoding is crucial to achieving photorealism , and significantly outperformed the previous state of the art , which included other coordinate-based MLPs ( Sitzmann et al. , 2019 ) . Concurrently , Zhong et al . ( 2020 ) demonstrated the usefulness of Fourier encoding for reconstruction of 3D protein complexes . Tancik et al . ( 2020 ) conducted an in-depth study of positional encoding for coordinate-based MLPs . They show that , under the Neural Tangent Kernel ( NTK ) theory ( Jacot et al. , 2018 ) , the positional encoding using a sinusoidal basis results in a stationary kernel , which in turn can be interpreted as a convolutional ( shift-invariant ) reconstruction filter , desirable for signals on Euclidean spaces . The theory also explains the “ spectral bias ” ( Rahaman et al. , 2019 ) of these models via the rapidly decaying NTK eigenvalues for MLPs . In practice , Tancik et al . ( 2020 ) propose randomly sampling the frequencies of the Fourier basis elements for the encoding , and show that it outperforms previous coordinate-based MLPs with no positional encoding ( Mescheder et al. , 2019 ) , or axis-aligned positional encoding ( Mildenhall et al. , 2020 ) . While groundbreaking results and insightful theoretical developments have been demonstrated , most of the attention so far has been focused on coordinate-based MLPs for functions on Euclidean spaces . In this paper , we focus on the non-Euclidean case . One alternative to Fourier features was introduced by Sitzmann et al . ( 2020 ) , who use the sine function as the MLP nonlinearity . Major differences are that the sinusoidal is applied on every layer , and the frequencies are defined by the network MLP weights . Zheng et al . ( 2021 ) studied the trade-off between memorization and generalization in positional encoding , showing the former is related to the rank of the embedding matrix and the latter is related to the distance preservation of the embedding . They also propose another alternative to Fourier features that consists of sampling a Gaussian at fixed offsets from the input . 3 BACKGROUND . NeRF ( Mildenhall et al. , 2020 ) recently demonstrated the importance of positional encoding for coordinate-based learning of high frequency functions . For a coordinate x = [ x1 , x2 , x3 ] ∈ R3 , the following map was applied before the multilayer perceptron ( MLP ) , xi 7→ { sin ( 20πxi ) , cos ( 20πxi ) , sin ( 21πxi ) , cos ( 21πxi ) , · · · , sin ( 2Lπxi ) , cos ( 2Lπxi ) } . ( 1 ) Recall that { sin ( m⊤x ) } ∪ { cos ( n⊤x ) } with m , n ∈ Z3 form an orthonormal basis for functions defined on a compact subset of R3 . So Eq . ( 1 ) corresponds to evaluating basis functions of different frequencies at input coordinates ; we will refer to this method as “ Euclidean encoding ” . In the particular case of Eq . ( 1 ) , only axis-aligned , powers-of-two frequencies are used , so the basis for bandwidth 2L is only sparsely sampled . Tancik et al . ( 2020 ) found it better to randomly sample the frequency space , which results in basis elements that are not axis aligned , but it is still a sparse sample of the complete basis . In NeRF ( Mildenhall et al. , 2020 ) , a subset of the input encodes the view direction . A direction can be associated with a point on the surface of the unit sphere , a non-Euclidean manifold . In order to apply the Euclidean encoding to such input , ( Mildenhall et al. , 2020 ) represent the direction as a 3D unit vector instead , which increases the input dimensionality from 2D to 3D , and causes the positional encoding to sample from a set of O ( b3 ) basis functions , instead of the minimal O ( b2 ) for a bandwidth b . In IPDF ( Murphy et al. , 2021 ) , the input coordinates are flattened 3× 3 rotation matrices , which correspond to points on the 3D rotation manifold SO ( 3 ) . The Euclidean encoding then corresponds to sampling basis functions on R9 . This results in sampling from a space of O ( b9 ) basis functions instead of O ( b3 ) . Our key observation is that positional encoding of input coordinates should reflect the geometry of the underlying manifold . We propose to use the appropriate orthonormal basis functions on the manifold to implement this idea . While for many applications the Euclidean manifold is the correct one , for others it is not , as exemplified above . In the following sections we will show how the appropriate basis can be found for different types of manifolds , how to actually choose a subset of the basis to use in the encoding , and demonstrate experimentally the benefits of using them . 4 METHOD . 4.1 OVERVIEW . We consider a neural network that takes as inputs points on an n-dimensional manifold M. We assume there exists an orthonormal basis B = { bi : M → R } for the space of scalar functions on the manifold L2 ( M ) . Our goal is to apply the map γ : M → Rk to positionally encode x ∈ M before feeding it to the network , where k is the encoding dimension and γ ( x ) = { bj ( x ) | bj ∈ B′ ⊂ B } . ( 2 ) The first question that arises is how to find the basis for a given manifold . There is no general solution , but specific solutions exist for some large classes of manifolds . For compact groups ( not necessarily abelian ) , the Peter-Weyl theorem establishes a generalization of the Fourier series that results in a countable set of orthonormal basis functions . It gives rise to the Wigner-D matrix elements that are used as a basis for SO ( 3 ) in Section 4.3 . For locally compact abelian groups , the Pontryagin duality can be used , which is a generalization of the Fourier transform . For Riemannian manifolds , an orthonormal basis can be obtained as eigenfunctions of the Laplace-Beltrami operator . For the sphere S2 , these are the spherical harmonics used in Section 4.2 . The neural tangent kernel ( NTK ) theory ( Jacot et al. , 2018 ; Arora et al. , 2019 ; Lee et al. , 2019 ) shows that , under certain conditions , an MLP trained for regression f : Rn → R , converges to the solution of a kernel regression f ( x ) ≈ ∑ i wiyik ( x , xi ) , where k is the NTK , wi the kernel regression weights , and the sum is over the whole dataset . Moreover , when the inputs x have constant norm , the kernel depends only on inner products : k ( x , y ) = k′ ( x⊤y ) . In other words , it is rotationinvariant . Tancik et al . ( 2020 ) leverages this theory to explain the success of Euclidean encoding – it transforms the rotation-invariant kernel in a shift-invariant one , which results in approximating a convolution operation over the whole training set for inputs in a flat space . In this paper , the inputs do not lie on Euclidean space , so the shift-invariance in the sense of translations in Euclidean space is not appropriate . For the manifolds considered , we will seek positional encodings that have constant norm ( lie on Sn ) and such that the natural shift on the manifold also preserves the inner products between positionally encoded inputs . Then , the NTK should approximate a convolution on the manifold . We ’ ll show in the following sections that these properties constrain the selection of basis elements .

Summary: The paper generalizes the commonly used sinusoidal position encoding scheme (such as that in NerF input encoding) to inputs naturally residing in non-Euclidean manifold.  This is achieved by representing the input coordinates as  projections on  alternative sets of orthornormal bases instead of the trivial Euclidean coordinates.   The authors argue that by doing so, they extend the translational invariance of Fourier features in Euclidean space to non-Euclidian manifolds, resulting in a convolution-like operator that is invariant to the 'natural shift' on these manifolds.  These new position encoding techniques are tested for several tasks and seem to get consistent improvemnts over their Fourier counterparts.