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 proposes a new 3D shape representation learning method using multi-scale wavelet decomposition. In particular, the authors introduce a neural network architecture that decomposed 3D shapes into sub-bands components at multiple scales. In particular, starting from a pointcloud the proposed model learns to decompose it into coarse (high frequency) and detail (low frequency) components using an adaptive lifting scheme, similar to the original lifting scheme introduced for defining second-generation wavelets. Subsequently, two transformer models are used to refine the coarse and approximate geometry of the 3D shape. The proposed model achieves state-of-the-art results on the shape classification task on the ModelNet40 and the ScanObjectNN dataset and on the part segmentation task on the ShapeNet Part dataset. The concept of using such an adaptive lifting scheme seems to facilitate learning and to the best of my knowledge is novel for the case of shape representation learning.