Adaptive Wavelet Transformer Network for 3D Shape Representation Learning

1 INTRODUCTION Figure 1 : The detail component of table , car , chair , and cup from the test set of ModelNet40 Wu et al . ( 2015 ) . Colors represent the absolute values of the detail coefficients per point . The larger values ( redder colors ) indicate , the higher probability that the points locate on high-frequency parts , e.g. , edges or non-flat areas . Such geometric representations are learned implicitly in the classification task . Analysis of 3D point cloud is a crucial topic in computer vision and graphics , and has wide applications in robotics Chen et al . ( 2019 ) , autonomous driving Qi et al . ( 2018 ) , and visual SLAM Hitchcox & Forbes ( 2020 ) , etc . To better understand the 3D shapes , effective point cloud analysis techniques and methods are in great demand . With the thriving of deep learning , numerous welldesigned neural networks are applied to extract expressive semantics of 3D point clouds with layered operations , in contrast to low-level handcrafted shape descriptors . The pioneering work PointNet Qi et al . ( 2017a ) learns a spatial encoding of each point using multi-layer percep- trons and aggregates all point features to a global shape representation . The follow-on work PointNet++ Qi et al . ( 2017b ) adopts a hierarchical encoder-decoder structure to process local regions , which down-samples point clouds in layers and gradually interpolates them back to the original resolution . The following works learn 3D shape representation from different perspectives : Xu et al . ( 2018 ) ; Thomas et al . ( 2019 ) extend regular grid convolution to irregular 3D point clouds , Li et al . ( 2018 ) ; Wu et al . ( 2019 ) design local operations on point patches , Wang et al . ( 2019 ) ; Shi & Rajkumar ( 2020 ) treat point clouds as graphs and apply graph convolution networks , to name a few . An interesting yet challenging problem is how to decompose 3D shapes into interpretable representations and exploit them for better shape analysis . For instance , some attempts Wang et al . ( 2019 ) ; Lan et al . ( 2019 ) capture local geometric relations among center points and their neighbors . However , these works treat all points or local regions equally , making it hard to extract features of the most relevant and crucial geometric parts , e.g. , areas containing flat planes or areas containing edges . Another challenging problem is how to analyze 3D shapes at multiple scales/resolutions in an explainable way . Though some works Qi et al . ( 2017b ) ; Jiang et al . ( 2019 ) extract point cloud features hierarchically , these works build the hierarchy in spatial/semantic spaces based on pairwise point distance , which is less interpretable w.r.t . shape geometry . The geometric components with different frequencies in 3D shapes contain distinct geometric characteristics ( e.g. , high frequency components correspond to edges / sharp parts , while low frequency components correspond to flat areas / smooth parts ) and provide complementary geometric information to the shape representations . To effectively learn 3D shape representations , it is desirable to separate these components and process them differently , rather than feed the whole shapes into a single processor . It , therefore , motivates us to decompose 3D shapes into sub-bands components . In this work , we propose a principled method for 3D shape representation learning rooted in wavelet transform . Theoretical properties of wavelets in signal processing , such as time-frequency localization and multi-resolution analysis , have been well studied , making it appropriate to fuse into deep neural networks with interpretability . Inspired by that , we propose to perform multi-resolution analysis within neural networks through lifting scheme Sweldens ( 1998 ) to achieve a data-driven wavelet transform . The multi-resolution analysis generates decomposed visual representations at each scale , which contributes to the interpretability of the network . The generated approximation or detail coefficients per point at each scale , corresponding to high or low frequency components , capture essential geometric interest for down-streaming tasks . Specifically , we first adapt wavelet transform that analyzes 1D temporal or 2D image signals to 3D shapes at different semantic levels and decompose the original point cloud into the approximation/coarse and detail components , as explained in Section 4.1 and 4.2 . Next , in Section 4.3 , we utilize Transfomers Vaswani et al . ( 2017 ) ; Choromanski et al . ( 2021 ) to pay different attention to features from approximation and detailed variation components , and then fuse them with the original point features , respectively . As shown in Figure 1 , the approximation and detail components have distinct but complementary effects on reflecting the geometry of 3D shapes . Equipped with wavelet transform and Transformers , we propose Adaptive Wavelet Transformer Network ( AWTNet ) that captures and refines the holistic and complementary geometry of 3D shapes to supplement neighboring local information . AWT-Net is trained using the loss functions proposed in Section 4.4 . Experimental results on challenging benchmarks demonstrate that our AWT-Net achieves state-ofthe-art or competitive performance on 3D shape classification and segmentation tasks . Comprehensive analysis and visualizations verify that our model effectively extracts the complementary geometric features from two decomposed components at multiple scales . For example , as shown in Figure 1 , we observe that our model can learn meaningfully samplings of point clouds in different categories and can detect key regions consistently within a semantic class , e.g. , it separates chair legs and backs from chair seat , and separates cup handle from the body . 2 RELATED WORK . Geometric Point Cloud Models . PointNet Qi et al . ( 2017a ) and DeepSet Zaheer et al . ( 2017 ) are pioneering neural network models that directly process point cloud to learn a spatial encoding of each point and then aggregate all individual point features to a holistic representation . Point-Net++ Qi et al . ( 2017b ) partitions point cloud into overlapping local regions by the distance metric of spatial and/or semantic spaces and extracts local features capturing fine geometric structures from neighbors hierarchically . DGCNN Wang et al . ( 2019 ) reconstructs the k-NN graph using nearest neighbors in feature space to capture local geometric structure while maintaining permutation invariance . GeoCNN Lan et al . ( 2019 ) defines a convolution operation that aggregates edge features to capture the local geometric relations . Xu et al . ( 2020 ) aggregates points from local neighbors with similar semantics in both Euclidean and Eigenvalue space . RS-CNN Liu et al . ( 2019c ) extends regular CNN to encode geometric relation in a local point set by learning the topology . DensePoint Liu et al . ( 2019b ) recognizes the implicit geometry by extracting contextual shape semantics . Our model learns point cloud geometry from a novel perspective , i.e. , multi-resolution wavelet analysis which is extensively applied in signal processing but seldom explored for 3D shapes . Wavelet in Vision . Theoretical properties of multi-resolution analysis using wavelets have been well studied , making such approaches more interpretable than CNNs . Several prior works incorporate wavelet representations into CNNs . Oyallon et al . ( 2017 ) proposed a hybrid network that replaces the first layers of ResNet with a wavelet scattering network . This modified ResNet resulted in comparable performance to the original ResNet but has a smaller number of trainable parameters . The wavelet pooling layer proposed in Williams & Li ( 2018 ) subsamples image features based on second-level wavelet decomposition . Wavelet-SRNet Huang et al . ( 2017 ) reconstructs super-resolution images from multi-scale wavelet coefficients estimated from low-resolution images under coefficient supervision . Similar scheme is adapted in Liu et al . ( 2020 ) for image demoiréing and Rong et al . ( 2020 ) for burst denoising . In Ramamonjisoa et al . ( 2021 ) , they decompose a single image using Haar Transform and then reconstruct depth images hierarchically through inverse discrete wavelet transform . All of the mentioned works use pre-defined Haar wavelet which is less flexible for complex 3D shapes . Rustamov & Guibas ( 2013 ) adopts lifting scheme but still restricts to linear transformations of Haar coefficients for graph signals . DAWN-Net Rodriguez et al . ( 2020 ) uses lifting scheme to implicitly define wavelet transform functions for image recognition . Applying lifting scheme to learn 3D shape representations is more challenging as amounts of shape points are arranged irregularly , and as far as we know , we are among the first to tackle this problem . Transformer in Vision . The seminal work to adapt Transformer Vaswani et al . ( 2017 ) to vision tasks commences with Vision Transformer ( ViT ) Dosovitskiy et al . ( 2020 ) and Detection Transformer Carion et al . ( 2020 ) . Since then , various applications of Transformer in vision are proposed , including : semantic segmentation Wang et al . ( 2021 ) ; Zhang et al . ( 2020 ) , generative modeling Chen et al . ( 2020 ) ; Hudson & Zitnick ( 2021 ) , object detection Zhu et al . ( 2021 ) ; Dai et al . ( 2021 ) , video analysis Gabeur et al . ( 2020 ) ; Zeng et al . ( 2020 ) , to name a few . PCT Guo et al . ( 2021 ) applies Transformer in 3D point cloud classification and segmentation . Concurrent works include two types of Point Transfomer proposed by Zhao et al . ( 2021 ) and Engel et al . ( 2021 ) . GDA-Net is built on graph spectral theory and uses a 2-order polynomial approximation ( Laplacian ) as the filter and applies it to a graph to disentangle the graph into high and low frequency components . Thus , the decomposition process in is deterministic and static . PoinTr Yu et al . ( 2021 ) converts a point cloud to a sequence of point proxies and employs Transformer for point cloud completion . PPT-Net Hui et al . ( 2021 ) adopts Transformer for large-scale place recognition . MViT Fan et al . ( 2021 ) and Swin Liu et al . ( 2021 ) connect multi-scale feature hierarchies with transformer models by applying different Transformers for features with distinct spatial resolutions . Since multi-scale discrete wavelet transform decomposes point clouds at each subsequent scales , it is intrinsically essential to combine with Transformer in a similar way as in MViT and Swin . Due to quadratic computational complexity of self-attention , Fast Autoregressive Transformers Katharopoulos et al . ( 2020 ) is among the first to address the quadratic complexity by using the associativity property of matrix products . Some recent works Kitaev et al . ( 2020 ) ; Dai et al . ( 2020 ) ; Tay et al . ( 2020 ) ; Choromanski et al . ( 2021 ) propose efficient Transformers for long sequences . These models can be also adapted for vision tasks and the one proposed in Choromanski et al . ( 2021 ) is explored in our work . 3 BACKGROUND . We briefly review wavelet transform and lifting scheme as the building blocks of our model . We refer the reader to Sweldens ( 1996 ) ; Claypoole et al . ( 2003 ) ; Rodriguez et al . ( 2020 ) for more information . Wavelet Transform/Decomposition . Let us consider a continuous scalar function of a single variable f ( x ) . Wavelet transform is to find a basis generated by a mother wavelet , denoted by ψ ( x ) , to represent f ( x ) with a few parameters . The basic rule for generating this basis is to translate and dilate ψ ( x ) , i.e. , ψ ( 2jx− i ) , where j and i are integers . The basis is localized in both time ( controlled by i ) and scale/frequency ( controlled by j ) domains . The original function f ( x ) is decomposed as a linear combination of the basis : f ( x ) = ∑ j , i aj , iψ ( 2 jx− i ) , ( 1 ) where aj , i are the wavelet coefficients of f ( x ) and can be computed as the inner product of ψ ( 2jx − i ) with f ( x ) . For a discretized signal x [ n ] ( n ∈ Z ) , we sample points from ψ ( 2jx − i ) to form filter banks to decompose x [ n ] . This is classical or first-generation wavelet Santosa ( 2011 ) . Some commonly used mother wavelets include Haar wavelet Haar ( 1910 ) , Daubechies wavelet , Biorthogonal wavelet Daubechies ( 1992 ) , to name a few . Wavelet transform provides a nice framework for multi-resolution analysis of functions , i.e. , the wavelet transform generates coarser and coarser versions of the same function while keeps track of the details missed going from fine to coarse . However , the first-generation wavelet mainly works well for infinite or periodic signals , which are rare in many applications . It is still unclear how to design a proper basis for any arbitrary signals sampled irregularly from a bounded domain . One solution is to give up designing a mother wavelet manually . Instead , the decomposition is entirely based on an approach called lifting scheme . Lifting Scheme . Lifting scheme , a.k.a. , second-generation wavelets Sweldens ( 1998 ) ; Daubechies & Sweldens ( 1998 ) , is a simple yet effective approach to define wavelets sharing the same properties as the first-generation wavelets . The lifting scheme takes as input a signal x and generates as outputs the approximation c and the detail d sub-bands of the wavelet transform . The lifting scheme consists of three stages Claypoole et al . ( 2003 ) as on the left panel of Figure 2 . • Split . This stage splits the input signal into two non-overlapping partitions , e.g. , the input signal x is divided into even xe and odd xo components , which are defined as xe [ n ] = x [ 2n ] and xo [ n ] = x [ 2n+ 1 ] . • Update ( U ) . This stage separates the signal in frequency domain , generating the approximation c that has the same running average as the original input signal . Let xLUo [ n ] = xo [ n − LU ] , xo [ n − LU + 1 ] , · · · , xo [ n + LU − 1 ] , xo [ n + LU ] denote the sequence of 2LU + 1 neighboring odd polyphase samples around xe [ n ] . The even polyphase samples are updated through an update operator U ( · ) on xLUo [ n ] , resulting the approximation c as : c [ n ] = xe [ n ] + U ( x LU o [ n ] ) . ( 2 ) • Predict ( P ) . As the partitions xe and xo are usually closely correlated , given one of them , it is feasible to build a predictor P ( · ) for the other , by tracking the difference ( or detail ) d among them . As the even component xe corresponds to the approximation c [ n ] , it is possible to define P ( · ) as a function of c [ n ] . Let cLP [ n ] = c [ n−LP ] , c [ n−LP +1 ] , · · · , c [ n+ LP − 1 ] , c [ n + LP ] denote a sequence of 2LP + 1 approximation coefficients . The odd polyphase samples are predicted as P ( cLP [ n ] ) . The resulting prediction residual , or detail d , are computed as : d [ n ] = xo [ n ] − P ( cLP [ n ] ) . ( 3 ) Note that the update and predict operators U and P implicitly and effectively define the wavelet for decomposition .

Summary: This paper presents a novel framework for 3D shape representation learning, which is based on multi-scale wavelet decomposition. This is very different from existing works. A novel transformer-based neural network, AWT-Net, is also proposed.