summaries: This paper proposes a novel 3D point cloud representation learning framework. At the core of this method, is a lifting scheme inspired by wavelet decomposition. The proposed method roughly splits the input data in half at each stage, producing a down-sampled approximation C and detail d. Then C is further processed by the next layer, forming a multiscale pyramid. In summary, the contribution of this paper is:  1. Proposed to use the lifting scheme in point cloud processing, using graph convolution networks and transformers as backbone.  2. Evaluated the method against state-of-the-art baselines and showed that the proposed scheme performs well.

summaries: 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.

summaries: 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.

summaries: The paper presents a new deep neural network architecture for 3D point cloud representation learning, based on wavelet decomposition. In particular, the authors propose a data-driven adaptive lifting scheme that introduces non-linearity into wavelet. The original linear operators update(U) and predict(P) in wavelet decomposition are replaced by non-linear graph convolutional networks (GCN). Equipped with wavelet transform and Transformers, the proposed network aims to captures and refines the holistic and complementary geometry of 3D shapes to supplement neighboring local information. Experimental results on standard benchmarks (i.e., shape classification and part segmentation) show that it achieves state-of- the-arts or competitive performance.

abstractText: We present a novel method for 3D shape representation learning using multi-scale wavelet decomposition. Distinct from previous works that either decompose 3D shapes into complimentary components at a single scale, or naively adopt up/down-sampling to build hierarchies and treat all points or local regions equally, we decompose 3D shapes into sub-bands components at multiple scales and all scales form a decomposition tree in a principled manner rooted in multi-resolution wavelet analysis. Specifically, we propose Adaptive Wavelet Transformer Network (AWT-Net) that firstly generates approximation or detail wavelet coefficients per point, classifying each point into high or low sub-bands components, using lifting scheme at multiple scales recursively and hierarchically. Then, AWT-Net exploits Transformers that regard the features from different but complementary components as two holistic representations, and fuse them with the original shape features with different attentions. The wavelet coefficients can be learned without direct supervision on coefficients, and AWT-Net is fully differentiable and can be learned in an end-to-end fashion. Extensive experiments demonstrate that AWT-Net achieves state-of-the-art or competitive performance on 3D shape classification and segmentation benchmarks.

authors:
Nicolas Carion
Francisco Massa
Gabriel Synnaeve
Nicolas Usunier
Alexander Kirillov
Sergey Zagoruyko

title: End-to-end object detection with transformers

venue: In Proceedings of the European Conference on Computer Vision,

authors:
Jingdao Chen
Yong Kwon Cho
Zsolt Kira

title: Multi-view incremental segmentation of 3-d point clouds for mobile robots

venue: IEEE Robotics and Automation Letters,

authors:
Mark Chen
Alec Radford
Rewon Child
Jeffrey Wu
Heewoo Jun
David Luan
Ilya Sutskever

title: Generative pretraining from pixels

venue: In International Conference on Machine Learning,

authors:
Krzysztof Marcin Choromanski
Valerii Likhosherstov
David Dohan
Xingyou Song
Andreea Gane
Tamas Sarlos
Peter Hawkins
Jared Quincy Davis
Afroz Mohiuddin
Lukasz Kaiser
David Benjamin Belanger
Lucy J Colwell
Adrian Weller

title: Rethinking attention with performers

venue: In International Conference on Learning Representations,

authors:
Roger L Claypoole
Geoffrey M Davis
Wim Sweldens
Richard G Baraniuk

title: Nonlinear wavelet transforms for image coding via lifting

venue: IEEE Transactions on Image Processing,

authors:
Zhigang Dai
Bolun Cai
Yugeng Lin
Junying Chen

title: Up-detr: Unsupervised pre-training for object detection with transformers

venue: In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition,

authors:
Zihang Dai
Guokun Lai
Yiming Yang
Quoc Le

title: Funnel-transformer: Filtering out sequential redundancy for efficient language processing

venue: In Advances in Neural Information Processing Systems,

authors:
Ingrid Daubechies
Wim Sweldens

title: Factoring wavelet transforms into lifting steps

venue: Journal of Fourier Analysis and Applications,

authors:
Alexey Dosovitskiy
Lucas Beyer
Alexander Kolesnikov
Dirk Weissenborn
Xiaohua Zhai
Thomas Unterthiner
Mostafa Dehghani
Matthias Minderer
Georg Heigold
Sylvain Gelly

title: An image is worth 16x16 words: Transformers for image recognition at scale

authors:
Haoqi Fan
Bo Xiong
Karttikeya Mangalam
Yanghao Li
Zhicheng Yan
Jitendra Malik
Christoph Feichtenhofer

title: Multiscale vision transformers

venue: arXiv preprint arXiv:2104.11227,

authors:
Stephen Fitzpatrick
Lambert Meertens

title: An experimental assessment of a stochastic, anytime, decentralized, soft colourer for sparse graphs

venue: In International Symposium on Stochastic Algorithms,

authors:
Valentin Gabeur
Chen Sun
Karteek Alahari
Cordelia Schmid

title: Multi-modal transformer for video retrieval

authors:
Meng-Hao Guo
Jun-Xiong Cai
Zheng-Ning Liu
Tai-Jiang Mu
Ralph R Martin
Shi-Min Hu

title: Pct: Point cloud transformer

venue: Computational Visual Media,

title: Zur theorie der orthogonalen funktionensysteme

venue: Mathematische Annalen,

authors:
Thomas Hitchcox
James Richard Forbes

title: A point cloud registration pipeline using gaussian process regression for bathymetric slam

venue: IEEE International Conference on Intelligent Robots and Systems,

authors:
Huaibo Huang
Ran He
Zhenan Sun
Tieniu Tan

title: Wavelet-srnet: A wavelet-based cnn for multiscale face super resolution

venue: In Proceedings of the IEEE International Conference on Computer Vision,

authors:
Drew A Hudson
C Lawrence Zitnick

title: Generative adversarial transformers

venue: arXiv preprint arXiv:2103.01209,

authors:
Le Hui
Hang Yang
Mingmei Cheng
Jin Xie
Jian Yang

title: Pyramid point cloud transformer for large-scale place recognition

authors:
Li Jiang
Hengshuang Zhao
Shu Liu
Xiaoyong Shen
Chi-Wing Fu
Jiaya Jia

title: Hierarchical point-edge interaction network for point cloud semantic segmentation

authors:
Angelos Katharopoulos
Apoorv Vyas
Nikolaos Pappas
François Fleuret

title: Transformers are rnns: Fast autoregressive transformers with linear attention

authors:
Thomas N Kipf
Max Welling

title: Semi-supervised classification with graph convolutional networks

authors:
Nikita Kitaev
Lukasz Kaiser
Anselm Levskaya

title: Reformer: The efficient transformer

authors:
Shiyi Lan
Ruichi Yu
Gang Yu
Larry S Davis

title: Modeling local geometric structure of 3d point clouds using geo-cnn

authors:
Yangyan Li
Rui Bu
Mingchao Sun
Wei Wu
Xinhan Di
Baoquan Chen

title: Pointcnn: Convolution on x-transformed points

venue: Advances in Neural Information Processing Systems,

authors:
Yiqun Lin
Zizheng Yan
Haibin Huang
Dong Du
Ligang Liu
Shuguang Cui
Xiaoguang Han

title: Fpconv: Learning local flattening for point convolution

authors:
Zhi-Hao Lin
Sheng-Yu Huang
Yu-Chiang Frank Wang

title: Convolution in the cloud: Learning deformable kernels in 3d graph convolution networks for point cloud analysis

authors:
Lin Liu
Jianzhuang Liu
Shanxin Yuan
Gregory Slabaugh
Aleš Leonardis
Wengang Zhou
Qi Tian

title: Wavelet-based dual-branch network for image demoiréing

authors:
Xinhai Liu
Zhizhong Han
Yu-Shen Liu
Matthias Zwicker

title: Point2sequence: Learning the shape representation of 3d point clouds with an attention-based sequence to sequence network

venue: In Proceedings of the AAAI Conference on Artificial Intelligence,

authors:
Yongcheng Liu
Bin Fan
Gaofeng Meng
Jiwen Lu
Shiming Xiang
Chunhong Pan

title: Densepoint: Learning densely contextual representation for efficient point cloud processing

venue: In Proceedings of the IEEE International Conference on Computer Vision, pp. 5239–5248,

authors:
Yongcheng Liu
Bin Fan
Shiming Xiang
Chunhong Pan

title: Relation-shape convolutional neural network for point cloud analysis

authors:
Ze Liu
Yutong Lin
Yue Cao
Han Hu
Yixuan Wei
Zheng Zhang
Stephen Lin
Baining Guo

title: Swin transformer: Hierarchical vision transformer using shifted windows

venue: arXiv preprint arXiv:2103.14030,

authors:
Jiageng Mao
Xiaogang Wang
Hongsheng Li

title: Interpolated convolutional networks for 3d point cloud understanding

authors:
Sunil K Narang
Antonio Ortega

title: Lifting based wavelet transforms on graphs

venue: In Proceedings of Asia-Pacific Signal and Information Processing Association,

authors:
Edouard Oyallon
Eugene Belilovsky
Sergey Zagoruyko

title: Scaling the scattering transform: Deep hybrid networks

authors:
Adam Paszke
Sam Gross
Francisco Massa
Adam Lerer
James Bradbury
Gregory Chanan
Trevor Killeen
Zeming Lin
Natalia Gimelshein
Luca Antiga

title: Pytorch: An imperative style, high-performance deep learning library

authors:
Charles R Qi
Hao Su
Kaichun Mo
Leonidas J Guibas

title: Pointnet: Deep learning on point sets for 3d classification and segmentation

authors:
Charles R Qi
Wei Liu
Chenxia Wu
Hao Su
Leonidas J Guibas

title: Frustum pointnets for 3d object detection from rgb-d data

authors:
Charles Ruizhongtai Qi
Li Yi
Hao Su
Leonidas J Guibas

title: Pointnet++: Deep hierarchical feature learning on point sets in a metric space

authors:
Ali Rahimi
Benjamin Recht

title: Random features for large-scale kernel machines

authors:
Michael Ramamonjisoa
Michael Firman
Jamie Watson
Vincent Lepetit
Daniyar Turmukhambetov

title: Single image depth prediction with wavelet decomposition

authors:
Maria Ximena Bastidas Rodriguez
Adrien Gruson
Luisa Polania
Shin Fujieda
Flavio Prieto
Kohei Takayama
Toshiya Hachisuka

title: Deep adaptive wavelet network

venue: In Proceedings of the IEEE Winter Conference on Applications of Computer Vision,

authors:
Xuejian Rong
Denis Demandolx
Kevin Matzen
Priyam Chatterjee
Yingli Tian

title: Burst denoising via temporally shifted wavelet transforms

authors:
Raif Rustamov
Leonidas J Guibas

title: Wavelets on graphs via deep learning

title: Second generation wavelets: theory and application

venue: University of Minnesota, Institute for Mathematics and its Applications,

authors:
Yiru Shen
Chen Feng
Yaoqing Yang
Dong Tian

title: Mining point cloud local structures by kernel correlation and graph pooling

authors:
Weijing Shi
Raj Rajkumar

title: Point-gnn: Graph neural network for 3d object detection in a point cloud

venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,

title: Wavelets and the lifting scheme: A 5 minute tour

venue: Journal of Applied Mathematics and Mechanics,

title: The lifting scheme: A construction of second generation wavelets

venue: SIAM Journal on Mathematical Analysis,

authors:
Yi Tay
Dara Bahri
Liu Yang
Donald Metzler
Da-Cheng Juan

title: Sparse sinkhorn attention

authors:
Hugues Thomas
Charles R Qi
Jean-Emmanuel Deschaud
Beatriz Marcotegui
François Goulette
Leonidas J Guibas

title: Kpconv: Flexible and deformable convolution for point clouds

authors:
Mikaela Angelina Uy
Quang-Hieu Pham
Binh-Son Hua
Thanh Nguyen
Sai-Kit Yeung

title: Revisiting point cloud classification: A new benchmark dataset and classification model on real-world data

authors:
Ashish Vaswani
Noam Shazeer
Niki Parmar
Jakob Uszkoreit
Llion Jones
Aidan N Gomez
Łukasz Kaiser
Illia Polosukhin

title: Attention is all you need

authors:
Huiyu Wang
Yukun Zhu
Hartwig Adam
Alan Yuille
Liang-Chieh Chen

title: Max-deeplab: Endto-end panoptic segmentation with mask transformers

authors:
Yue Wang
Yongbin Sun
Ziwei Liu
Sanjay E Sarma
Michael M Bronstein
Justin M Solomon

title: Dynamic graph cnn for learning on point clouds

venue: Acm Transactions On Graphics,

authors:
Travis Williams
Robert Li

title: Wavelet pooling for convolutional neural networks

authors:
Wenxuan Wu
Zhongang Qi
Li Fuxin

title: Pointconv: Deep convolutional networks on 3d point clouds

authors:
Zhirong Wu
Shuran Song
Aditya Khosla
Fisher Yu
Linguang Zhang
Xiaoou Tang
Jianxiong Xiao

title: 3d shapenets: A deep representation for volumetric shapes

authors:
Saining Xie
Sainan Liu
Zeyu Chen
Zhuowen Tu

title: Attentional shapecontextnet for point cloud recognition

authors:
Mingye Xu
Zhipeng Zhou
Yu Qiao

title: Geometry sharing network for 3d point cloud classification and segmentation

authors:
Mutian Xu
Junhao Zhang
Zhipeng Zhou
Mingye Xu
Xiaojuan Qi
Yu Qiao

title: Learning geometry-disentangled representation for complementary understanding of 3d object point cloud

authors:
Yifan Xu
Tianqi Fan
Mingye Xu
Long Zeng
Yu Qiao

title: Spidercnn: Deep learning on point sets with parameterized convolutional filters

authors:
Li Yi
Vladimir G Kim
Duygu Ceylan
I-Chao Shen
Mengyan Yan
Hao Su
Cewu Lu
Qixing Huang
Alla Sheffer
Leonidas Guibas

title: A scalable active framework for region annotation in 3d shape collections

venue: ACM Transactions on Graphics,

authors:
Li Yi
Hao Su
Xingwen Guo
Leonidas J Guibas

title: Syncspeccnn: Synchronized spectral cnn for 3d shape segmentation

authors:
Xumin Yu
Yongming Rao
Ziyi Wang
Zuyan Liu
Jiwen Lu
Jie Zhou

title: Pointr: Diverse point cloud completion with geometry-aware transformers

authors:
Yanhong Zeng
Jianlong Fu
Hongyang Chao

title: Learning joint spatial-temporal transformations for video inpainting

authors:
Dong Zhang
Hanwang Zhang
Jinhui Tang
Meng Wang
Xiansheng Hua
Qianru Sun

title: Feature pyramid transformer

authors:
Xizhou Zhu
Weijie Su
Lewei Lu
Bin Li
Xiaogang Wang
Jifeng Dai

title: Deformable detr: Deformable transformers for end-to-end object detection

text:
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.

heading: 2 RELATED WORK

text:
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.

heading: 3 BACKGROUND

text:
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.

heading: 4 ADAPTIVE WAVELET TRANSFORMER NETWORK

text:
Overview. The update and predict operators in Eq. 2 and Eq. 3 are linear, which are incapable of learning geometry of 3D shapes. In this section, we propose a data-driven adaptive lifting scheme that introduces non-linearity into wavelet. Based on this scheme, we develop an Adaptive Wavelet Transformer Network, dubbed as AWT-Net, for 3D shape representation learning. The AWT-Net
generates shape representations hierarchically, and the decomposition tree is shown on the left panel of Figure 3. Here we show three scales where each scale corresponds to a separate lifting process as illustrated on the left panel of Figure 2. We omit the split, update and predict blocks for conciseness. At each scale, the approximation ci serves the coarse version of the input signal while the detail di maintains the detailed difference between ci and the input signal. The approximation ci is further decomposed by another lifting process at the next scale. The AWT-Net is built based on this decomposition tree but replaces the linear operators U and P with non-linear neural networks. Once trained and fixed, the update U and predict P operators can be regarded as two sub-band basis/filters to subdivide the point cloud into high-frequency and low-frequency components, in analogy to the basis in Eq. 1. Its unrolled version is shown on the right panel of Figure 3.
Each module in AWT-Net is described as follows.

heading: 4.1 EVEN-ODD SPLIT

text:
Graph Construction. Assume a point cloud containing N points associated with features xi ∈ Rd, i ∈ {1, 2, · · · , N}. Features can be 3D coordinates, normals and/or high-dimensional features. We construct a k-NN graph G = (V, Ãu, Ãw), where each point xi corresponds to a node vi ∈ V , Ãu and Ãw ∈ RN×N are unweighted and weighted adjacency matrices encoding point similarity in feature space. The unweighted and weighted edges connecting two nodes vi and vj are defined as:
Ãuij [, Ãwij ] = { 1 [, κ(∥xi − xj∥2)], xj ∈ N (xi) 0, otherwise , (4)
where κ(·) is an non-negative function, e.g., Gaussian function, to ensure that Ãw is a diagonally dominant matrix andN represents neighborhood. Following Xu et al. (2021), we normalize all edge weights in each row in Aw to remove the impacts of varying neighborhood sizes and feature scales. Graph Split. We partition all nodes in a graph into two sets, containing even and odd nodes, respectively. Since lifting scheme uses information from odd (resp. even) nodes to update (resp. predict) even (resp. odd) nodes as in Eq. 2 (resp. Eq. 3), having all neighboring nodes with same parity means that local information cannot be used. Following Narang & Ortega (2009), we define the number of conflicts as the percentage of neighboring nodes that have the same parity after partition. The optimal splitting minimizes the number of conflicts. Theoretically, it is an NP-hard bipartite subgraph problem, but a heuristic solution can be found. We adapt the conservative fixed probability colorer (CFP) algorithm in Fitzpatrick & Meertens (2001), which is based on iteratively greedy local heuristics: at each iteration, a few randomly chosen nodes are activated, and each activated node counts the number of conflicting edges (i.e., connecting neighboring nodes with the same parity) and changes its parity based on the conflict. The algorithm is summarized in Algorithm 1.
Algorithm 1: Even-Odd Node Split Require: A graph G = (V,Au) with unweighted adjacency, maximum number of iterations M
1 Randomly assign initial parity (e.g., 1 for odd, 0 for even) to each node 2 for m = 1:M do 3 Activate each node with a probability sampled from a uniform distribution independently 4 for each activated node do 5 Choose a parity that minimizes its conflict with neighboring nodes
/* Keep the number of odd and even nodes roughly the same */ 6 No, Ne ← number of nodes with parity1, parity0 7 if No > Ne then 8 Randomly flip ⌊No−Ne2 ⌋ odd nodes’ parity 9 else
10 Randomly flip ⌊Ne−No2 ⌋ even nodes’ parity

heading: 4.2 ADAPTIVE LIFTING SCHEME

text:
Based on nodes’ parity after the even-odd split, we rearrange nodes V to gather odd and even nodes, and also rearrange the weighted adjacency matrix Aw accordingly:
V = ( Vodd Veven ) , Aw = ( FNo×No JNo×Ne KNe×No LNe×Ne ) ,
where Vodd and Veven represent the sets of odd and even nodes. The sub-matrix F is the adjacency matrix of a subgraph containing odd nodes only. The submatrix J contains edges that do not have conflicts, i.e., edges connecting nodes with the opposite parity. Similar definitions for K and L. Unlike the linear lifting scheme introduced in Section 3, we propose data-dependent non-linear update and predict operators, implemented as neural networks whose parameters are adaptively optimized to capture complex shape geometry. Specifically, we apply two Graph Convolutional Networks (GCNs) proposed in Kipf & Welling (2017) as the update and predict operators, separately. The graph convolution allows us to compute the response of a node based on its neighbors decided by the adjacency matrix. Performing graph convolution is equal to performing message gather within a neighborhood around each center node, sharing similarities with the definitions of xLuo [n] and cLp [n] in Eq. 2 and Eq. 3. The non-linear update and predict stages at i-th scale are defined as:
ci = Veven + GCN(Vodd,F) , (5) di = Vodd − GCN(ci,L) , (6)
where ci and di represent the approximation and detail components at i-th scale. We repeat this process at each scale recursively and form a hierarchical decomposition tree, capturing shape geometry at multi-scales. The maximum number of scales can be decided analytically as log2N and the coarsest scale contains only a single point.

heading: 4.3 TRANSFORMER FOR CROSS ATTENTION

text:
The approximation and detail components play different but complementary roles in representing 3D shape geometry, and thus should not be treated equally. To utilize different variation components gained from lifting scheme, we adopt Transformers Vaswani et al. (2017); Choromanski et al. (2021) to integrate different components with the original input shape features. The architecture of Transformers is shown on the right panel of Figure 2, and multi-head attention is defined as:
Attention(Q,K, V ) = softmax(QK⊤)V , (7)
where Q,K, V denote query, key and value, respectively.
For approximation ci as coarse component at i-th scale, we set query, key and value to:
Q = zi−1,K = ci, V = ci , (8)
where the definition of zi is shown in Figure 3. Note that for the first scale (i.e., i = 1), we set Q = x as input. For detail di, as most of the values are close to zeros (explained in Subsection 4.4), directly feed them into Transformer yields sub-optimal results. Utilizing a nice property of wavelets, namely time-frequency localization1, we regard di as a multiplier to strengthen or weaken the context features. Formally, the query, key and value for detail di at i-th scale are set as:
Q = zi−1,K = ci−1 ⊗ exp(di), V = ci−1 ⊗ exp(di) , (9)
where ⊗ represents element-wise multiplication. For the first scale, we set both zi−1 and ci−1 to x. In the part segmentation task, since each point cloud contains 2048 points, to efficiently compute multi-head attention, we also test Performers Choromanski et al. (2021) in addition to the vanilla Transformers Vaswani et al. (2017). Briefly, Performers represent the (attention) softmax kernel in terms of the Gaussian kernel, and approximate the latter by sampling from the Fourier Transform Rahimi et al. (2007).

heading: 4.4 LOSS FUNCTION

text:
End-to-end training is performed using the cross-entropy loss function, in combination with regularization terms, to enforce a wavelet decomposition structure during training. The loss function takes the form of Eq. 10, where K in the first term denotes the number of classes, yi and pi are the ground-truth and the predicted probability for belonging to class i, respectively.
Loss =− K∑ i=1 yi log(pi) + λ1 S∑ i=1 SmoothL1(A2ci−1 − ci) + λ2 S∑ i=1 SmoothL1(di) . (10)
In lifting scheme, a perfect output of Eq. 2, i.e., the series c[n], would be a coarse approximation of x[n] at half the resolution. In the second term, we enforce each lifted even node has the same local average as its corresponding node in the original graph. Due to the half resolution of ci compared with the input ci−1, we raise the adjacent matrix A to the power of two to capture the node features up to second-order neighbors, and minimize the difference between the input graph features and the coarser version across all the decomposition scales. The rationale for the third term is: if the underlying signal x[n] is locally smooth, the prediction residuals d[n] should be small. Thus, the third term promotes the minimization of the magnitude of detail coefficients d[n] in Eq. 3, indicating smaller energy is contained in d[n].

heading: 5 EXPERIMENTS AND RESULTS

text: We evaluate our model on 3D shape classification and part segmentation tasks on different datasets.

heading: 5.1 IMPLEMENTATION DETAILS

text: The number of the constructed neighborhoods at each stage is 64. The maximum number of iterations M used in the split Algorithm 1 is set to 10. The update U(·) and P (·) operators are implemented as two GCNs Kipf & Welling (2017) and the internal structure of each GCN is: ConvBN-ReLU-Dropout-Conv-BN. We use a two-heads attention in Transformer (for both the vanilla one and Performer). We implement our model in PyTorch Paszke et al. (2019). We use the SGD optimizer with momentum and weight decay set to 0.9 and 0.0001, respectively. For both 3D shape classification and 3D object part segmentation tasks, we train for 350 epochs. The initial learning rate is set to 0.1 and is dropped by 5x every 40 epochs and clipped at 0.9e− 5. The values of λ1 and λ2 are set to 0.1 for all tasks.

heading: 5.2 OBJECT CLASSIFICATION

text:
We evaluate AWT-Net on ModelNet40 Wu et al. (2015) which contains 9843 training shapes and 2468 test shapes across 40 categories. Points are uniformly sampled from the mesh models by Qi
1In analogy wave signals expanding time domain, for 3D shapes which expand in the spatial domain, this property can be rephrased as spatial-frequency localization.
et al. (2017a). Following Wang et al. (2019); Xu et al. (2021), we translate shapes and then randomly shuffle points in the training set. Similar to Qi et al. (2017a;b), we perform several voting tests with random scaling and average the predictions during testing. The quantitative comparisons of the overall classification accuracy with the state-of-the-art methods are listed in Table 1. ‘+V’ denotes the vanilla Transformer Vaswani et al. (2017) and ‘+P’ denotes Performer Choromanski et al. (2021), a variant of the vanilla Transformer. AWT-Net with two scales achieves better or comparable performance with other methods using the same number of points as the input, as shown in Table 1.
We also evaluate AWT-Net on ScanObjectNN by Uy et al. (2019) to examine the robustness to deformed and noisy objects scanned in the real world. We test on the OBJ ONLY and noisy OBJ BG splits, and the results are summarized in Table 2, where our model gets the comparable accuracy and performance drop from OBJ ONLY to OBJ BG.
Method (time order) Points Accuracy
PointNet Qi et al. (2017a) 1K 89.2 PointNet++ Qi et al. (2017b) 1K 90.7 SCN Xie et al. (2018) 1K 90.0 KCNet Shen et al. (2018) 1K 91.0 PointCNN Li et al. (2018) 1K 92.2 Point2Sequence Liu et al. (2019a) 1K 92.6 DGCNN Wang et al. (2019) 1K 92.9 InterpCNN Mao et al. (2019) 1K 93.0 RS-CNN Liu et al. (2019c) 1K 93.6 3D-GCN Lin et al. (2020b) 1K 92.1 FPConv Lin et al. (2020a) 1K 92.5 GSNet Xu et al. (2020) 1K 92.9 PointTransformer Zhao et al. (2021) 1K 93.7 PCT Guo et al. (2021) 1K 93.2 GDANet Xu et al. (2021) 1K 93.8
AWT-Net +V (ours) 1K 93.5 AWT-Net +P (ours) 1K 93.9
Table 1: Classification accuracy (%) on ModelNet40 dataset.
Method OBJ ONLY OBJ BG Acc. drop
PointNet Qi et al. (2017a) 79.2 73.3 ↓ 5.9 PointNet++ Qi et al. (2017b) 84.3 82.3 ↓ 2.0 SpiderCNN Xu et al. (2018) 79.5 77.1 ↓ 5.4 PointCNN Li et al. (2018) 87.9 85.8 ↓ 2.1 DGCNN Wang et al. (2019) 86.2 82.8 ↓ 3.4 GDANet Xu et al. (2021) 88.5 87.0 ↓ 1.5 AWT-Net +V (ours) 88.2 86.4 ↓ 1.8 AWT-Net +P (ours) 88.5 86.8 ↓ 1.7
Table 2: Classification accuracy (%) w.r.t. noise on ScanObjectNN dataset.

heading: 5.3 PART SEGMENTATION

text: We evaluate AWT-Net with three scales for shape part segmentation on ShapeNet Part Yi et al. (2016) dataset, which contains 16881 shapes in 16 categories and is labeled in 50 parts where each shape has 25̃ parts. The same voting test in classification is adopted. The instance average, the class average, and each class mean Inter-over-Union (mIoU) are summarized in Table 3. AWTNet +P achieves the best performance with class mIoU of 85.0% and instance mIoU of 86.6%. It is worth noticing that AWT-Net performs better on objects with obvious boundaries such as bag, knife, laptop, and skateboard. Some qualitative segmentation results are visualized in Figure 4.

heading: 6 ABLATION STUDY

text:
Experiments are conducted using Performer (‘+P’) and no voting during testing (unless specified).
Module Analysis. We ablate each module in AWT-Net, and the results are listed in Table 4. When the Transformer takes as input the original point cloud and both approximation and detail components, the network achieves the best accuracy with 93.4%. AWT-Net also surpasses the architecture of only using k-NN with 1.1% improvement. Our method ultimately obtains an accuracy of 93.9% with voting tests. This experiment shows that the decomposed approximation and detail components capture different but complementary geometric features to get complete shape representation.
Split Scheme. We replace the split of approximation and detail components with random and farthest point sampling (FPS). The results are listed in Table 5, where decomposing point clouds into approximation and detail components gets the highest accuracy. Empirically, our wavelet decomposition strategy localizes points carried with informative geometry instead of noisy points, which improves the robustness against noise.
Adjacency Matrix Computation. The result of dynamic adjacency matrix calculation in Subsection 4.1 is shown in Table 6. We pre-compute the adjacency matrix on 3D coordinates to split point clouds, which gets the accuracy of 91.8%. Nevertheless, by dynamically calculating the adjacency matrix in high dimensional semantic space at different scales, our network gains 1.6% improvement. Figure 1 shows that our decomposition strategy successfully decomposes points into flat areas and areas containing edges. Note that we do not provide any supervision w.r.t. shape geometry.
Decomposition Levels. We investigate the impact of the number of decomposition levels (i.e., resolution scales), and the results are summarized in Table 7. AWT-Net performs best on the overall accuracy when the decomposition level is set to 2. We also observe that for shapes with simple geometric structures (e.g., radio and wardrobe), more scales have a negative effect on the performance, as more scales make the network prone to over-fitting on these simple shapes. While for shapes with complex structures (e.g., stool and person), more scales are needed to capture both coarse and fine-grained geometric features. That is, the accuracy increases with more scales applied to geometrically more complex objects. We find out that the ModelNet40 dataset contains more geometrically simple objects than complex ones, which leads to the overall accuracy drops with more scales used.
Robustness Analysis. First, the robustness of our network on sampling density is tested by using sparser points as the input to ATW-Net trained with 1024 points. Figure 5 lists the results. Although sparser points bring difficulties, ATW-Net still performs consistently across all density settings. Moreover, Table 8 summarizes the results of rotation robustness, where AWT-Net outperforms other compared methods, especially with 2.4% improvement than the second-best at (s/s). Last, our model is evaluated on ScanObjectNN Uy et al. (2019) that consists of noisy and deformed scanned objects with non-uniform surface density. Table 2 summarizes AWT-Net gains the comparable accuracy drop from OBJ ONLY to OBJ BG, proving that ATW-Net is robust against noise.

heading: 7 CONCLUSION

text: We propose AWT-Net for 3D shape representation learning. Rooted in wavelet transform, AWT-Net dynamically decomposes point clouds into approximation and detail components at multiple resolution scales, which respectively denotes the low and high sub-bands of point clouds recursively. Another core component is Transformer, which explores the relations between the original complete points and coarse and detail components to provide complementary geometric information for learning shape representations. Extensive experiments show that our method achieves state-of-the-art or comparable performances and is robust to noise.

heading: A APPENDIX

text:
A.1 FIGURES AND TABLES FOR ABLATION STUDY
Method ModelNet40 OBJ ONLY OBJ BG
random 89.6 82.9 83.8 FPS 92.0 85.3 84.2 app.-det. 93.4 88.0 86.3
Table 6: Impact of dynamic strategy.
Number of scales 1 2 3 4
Overall 92.5 93.4 93.1 91.3 Radio 90.0 83.5 75.0 75.0 Wardrobe 75.0 71.0 68.0 65.0 Stool 75.0 78.0 85.0 85.2
Person 85.0 88.5 95.0 100.0
Table 7: Effect of decomposition levels on ModelNet40. Note that the results are achieved without voting during inference.
Method z/z s/s
PointNet Qi et al. (2017a) 81.6 66.3 PointNet++ Qi et al. (2017b) 90.1 87.8 SpiderCNN Xu et al. (2018) 83.5 69.6 DGCNN Wang et al. (2019) 90.4 82.6
GDANet Xu et al. (2021) 91.2 90.5
AWT-Net (ours) 91.4 90.2
Table 8: Accuracy comparisons of rotation on ModelNet40. z/z: both training and test samples are augmented by random rotation along the z axis; s/s: both training and test samples are augmented by random rotation along (x,y,z) three axis.
A.2 EXECUTION TIME AND MODEL PARAMETERS
In this subsection, we compare the number of parameters and computational overhead w.r.t. GDANet Xu et al. (2021) on ModelNet40 and the results are listed in Table 9. We measure four scales. Note that GDA-Net consists of two building blocks, which is fairly comparable to “scale 2” in our model. In this case, though our model contains around 25% more parameters, the execution time is only 13% (+P) and 15% (+V) slower than GDA-Net, which shows the practicability of our model. We also notice that Performer is faster than the vanilla Transformer, which is consistent with our motivation.
A.3 RANDOMNESS IN EVEN-ODD NODE SPLIT
In this subsection, we compare the effect of using even-odd split and random split in Algorithm 1 and the results are listed in Table 11. Using random split performers slightly worse than even-odd split. With even-odd split, we can separate any two connected nodes into two different sets to the best extent, thus ensuring each local area contains both even and odd nodes. However, with random split, some local areas might contain only one type of nodes, thus failing to be decomposed properly.
A.4 CONFIGURATIONS OF MODULE ANALYSIS
In Table 11, we provide a description of the configurations of module analysis listed in Table 4.


summaries: This paper proposes a simple yet effective method, cocktail fine-tuning, for the natural language generation tasks. Their results show that cocktail fine-tuning can handle both In-domain data and Out-of-domain data effectively by combing adapter-finetuning and full-finetuning through knowledge distillation and overall has comparable performance compared to their ensembles. It also provides theoretical analysis on multi-class logistic regression to explain why it works.

summaries: This paper proposes an ensemble model between a full fine-tuning model and a parameter-efficient fine-tuning model to improve the out-of-distribution (OOD) performance of a full fine-tuning model. The proposed method is inspired by the observation that full fine-tuning model achieves good in-distribution (ID) performance while parameter-efficient finetuning model achieves better OOD performance. There are two ensembling methods presented in the paper: linear interpolation between the predictions of the two models; and distill from the predictions of a parameter-efficient model with ID training data. Improved OOD performance is observed with this ensemble method.

summaries: This paper presents interesting an idea of combining lightweight fine-tuning and full fine-tuning to achieve the best of both approaches, i.e. perform best on out-of-domain and in-domain data. The authors proposed two approaches: a simple ensemble method and a so-called cocktail fine-tuning that combines two fine-tuning methods in one single model. They evaluated their tasks in three datasets: WebNLG, XSUM and OpenQA and obtained mixed results. The authors also provided good analyses for more insights.

summaries: The present paper first discusses the trade-off between performance for out-of-domain data and in-domain data with respect to whether the model is fully fine-tuned or lightweight fine-tuned on NLG tasks. Second, it argues that such a trade-off is not necessary if one can make use of both of these two fine-tuning schema in a clever way. To this end, it proposes cocktail fine-tuning, which augments full fine-tuning via distillation from a lightweight model and which achieves equal performance as an ensemble of the two fine-tuning schema. At length, this paper also explains the behavior of the cocktail fine-tuning through a toy model.

abstractText: When finetuning a pretrained language model for natural language generation tasks, one is currently faced with a tradeoff. Lightweight finetuning (e.g., prefixtuning, adapters), which freezes all or most of the parameters of the pretrained model, has been shown to achieve stronger out-of-distribution (OOD) performance than full finetuning, which tunes all of the parameters. However, lightweight finetuning can underperform full finetuning in-distribution (ID). In this work, we present methods to combine the benefits of full and lightweight finetuning, achieving strong performance both ID and OOD. First, we show that an ensemble of the lightweight and full finetuning models achieves the best of both worlds: performance matching the better of full and lightweight finetuning, both ID and OOD. Second, we show that we can achieve similar improvements using a single model instead of two with our proposed cocktail finetuning, which augments full finetuning via distillation from a lightweight model. Finally, we provide some explanatory theory in a multiclass logistic regression setting with a large number of classes, describing how distillation on ID data can transfer the OOD behavior of one model to another.

authors:
Samira Abnar
Mostafa Dehghani
Willem Zuidema

title: Transferring inductive biases through knowledge distillation

venue: arXiv preprint arXiv:2006.00555,

authors:
Jonathan Berant
Andrew Chou
Roy Frostig
Percy Liang

title: Semantic parsing on Freebase from question-answer pairs

venue: In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing,

authors:
David Berthelot
Rebecca Roelofs
Kihyuk Sohn
Nicholas Carlini
Alex Kurakin

title: Adamatch: A unified approach to semi-supervised learning and domain adaptation

authors:
Cristian Bucilă
Rich Caruana
Alexandru Niculescu-Mizil

title: Model compression

venue: In International Conference on Knowledge Discovery and Data Mining (KDD),

authors:
Yair Carmon
Aditi Raghunathan
Ludwig Schmidt
Percy Liang
John C. Duchi

title: Unlabeled data improves adversarial robustness

venue: In Advances in Neural Information Processing Systems (NeurIPS),

authors:
Tommaso Furlanello
Zachary Lipton
Michael Tschannen
Laurent Itti
Anima Anandkumar

title: Born again neural networks

venue: In International Conference on Machine Learning,

authors:
Claire Gardent
Anastasia Shimorina
Shashi Narayan
Laura Perez-Beltrachini

title: The WebNLG challenge: Generating text from RDF data

venue: In Proceedings of the 10th International Conference on Natural Language Generation,

authors:
Geoffrey Hinton
Oriol Vinyals
Jeffrey Dean

title: Distilling the knowledge in a neural network

venue: In NIPS Deep Learning and Representation Learning Workshop,

authors:
Neil Houlsby
Andrei Giurgiu
Stanislaw Jastrzebski
Bruna Morrone
Quentin De Laroussilhe
Andrea Gesmundo
Mona Attariyan
Sylvain Gelly

title: Parameter-efficient transfer learning for NLP

venue: Proceedings of the 36th International Conference on Machine Learning,

authors:
Neil Houlsby
Andrei Giurgiu
Stanislaw Jastrzebski
Bruna Morrone
Quentin de Laroussilhe
Andrea Gesmundo
Mona Attariyan
Sylvain Gelly

title: Parameter-efficient transfer learning for NLP. arXiv, 2019b

authors:
Jeremy Howard
Sebastian Ruder

title: Universal language model fine-tuning for text classification. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)

authors:
Mihir Kale
Abhinav Rastogi

title: Text-to-text pre-training for data-to-text tasks

venue: In Proceedings of the 13th International Conference on Natural Language Generation,

authors:
Jungo Kasai
Hao Peng
Yizhe Zhang
Dani Yogatama
Gabriel Ilharco
Nikolaos Pappas
Yi Mao
Weizhu Chen
Noah A Smith

title: Finetuning pretrained transformers into rnns

venue: arXiv preprint arXiv:2103.13076,

authors:
Fereshte Khani
Percy Liang

title: Removing spurious features can hurt accuracy and affect groups disproportionately

venue: In ACM Conference on Fairness, Accountability, and Transparency (FAccT),

authors:
Tom Kwiatkowski
Jennimaria Palomaki
Olivia Redfield
Michael Collins
Ankur Parikh
Chris Alberti
Danielle Epstein
Illia Polosukhin
Matthew Kelcey
Jacob Devlin
Kenton Lee
Kristina N. Toutanova
Llion Jones
Ming-Wei Chang
Andrew Dai
Jakob Uszkoreit
Quoc Le
Slav Petrov

title: Natural questions: a benchmark for question answering research

venue: Transactions of the Association of Computational Linguistics,

authors:
Kenton Lee
Ming-Wei Chang
Kristina Toutanova

title: Latent retrieval for weakly supervised open domain question answering

venue: In Association for Computational Linguistics (ACL),

authors:
Lisa Lee
Benjamin Eysenbach
Emilio Parisotto
Eric Xing
Sergey Levine
Ruslan Salakhutdinov

title: Efficient exploration via state marginal matching

authors:
Brian Lester
Rami Al-Rfou
Noah Constant

title: The power of scale for parameter-efficient prompt tuning

authors:
Mike Lewis
Yinhan Liu
Naman Goyal
Marjan Ghazvininejad
Abdelrahman Mohamed
Omer Levy
Veselin Stoyanov
Luke Zettlemoyer

title: BART: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension

venue: In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics,

authors:
Xiang Lisa Li
Percy Liang

title: Prefix-tuning: Optimizing continuous prompts for generation

venue: Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pp. 4582–4597,

title: ROUGE: A package for automatic evaluation of summaries

venue: In Text Summarization Branches Out,

authors:
Pengfei Liu
Weizhe Yuan
Jinlan Fu
Zhengbao Jiang
Hiroaki Hayashi
Graham Neubig

title: Pre-train, prompt, and predict: A systematic survey of prompting methods in natural language processing

authors:
Robert L Logan IV
Ivana Balažević
Eric Wallace
Fabio Petroni
Sameer Singh
Sebastian Riedel

title: Cutting down on prompts and parameters: Simple few-shot learning with language models

venue: arXiv preprint arXiv:2106.13353,

authors:
Amir Najafi
Shin ichi Maeda
Masanori Koyama
Takeru Miyato

title: Robustness to adversarial perturbations in learning from incomplete data

authors:
Shashi Narayan
Shay B. Cohen
Mirella Lapata

title: Don’t give me the details, just the summary! Topic-aware convolutional neural networks for extreme summarization

venue: In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, Brussels,

authors:
Deepak Narayanan
Keshav Santhanam
Amar Phanishayee
Matei Zaharia

title: Accelerating Deep Learning Workloads through Efficient Multi-Model Execution

venue: In NeurIPS Workshop on Systems for Machine Learning,

authors:
Kishore Papineni
Salim Roukos
Todd Ward
Wei-Jing Zhu

title: Bleu: A method for automatic evaluation of machine translation

venue: In Proceedings of the 40th Annual Meeting on Association for Computational Linguistics,

title: BLEU: A method for automatic evaluation of machine translation

authors:
Jonas Pfeiffer
Aishwarya Kamath
Andreas Rücklé
Kyunghyun Cho
Iryna Gurevych

title: Adapterfusion: Non-destructive task composition for transfer learning, 2020

authors:
Guanghui Qin
Jason Eisner

title: Learning how to ask: Querying LMs with mixtures of soft prompts

authors:
Alec Radford
Jong Wook Kim
Chris Hallacy
Aditya Ramesh
Gabriel Goh
Sandhini Agarwal
Girish Sastry
Amanda Askell
Pamela Mishkin
Jack Clark
Gretchen Krueger
Ilya Sutskever

title: Learning transferable visual models from natural language supervision, 2021

authors:
Aditi Raghunathan
Sang Michael Xie
Fanny Yang
John C. Duchi
Percy Liang

title: Understanding and mitigating the tradeoff between robustness and accuracy

venue: In International Conference on Machine Learning (ICML),

authors:
Pranav Rajpurkar
Jian Zhang
Konstantin Lopyrev
Percy Liang

title: SQuAD: 100,000+ questions for machine comprehension of text

venue: In Empirical Methods in Natural Language Processing (EMNLP),

authors:
Sylvestre-Alvise Rebuffi
Hakan Bilen
Andrea Vedaldi

title: Learning multiple visual domains with residual adapters

authors:
Adam Roberts
Colin Raffel
Noam Shazeer

title: How much knowledge can you pack into the parameters of a language model

venue: In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP),

authors:
Taylor Shin
Yasaman Razeghi
Robert L. Logan IV au
Eric Wallace
Sameer Singh

title: Autoprompt: Eliciting knowledge from language models with automatically generated prompts, 2020

authors:
Jonathan Uesato
Jean-Baptiste Alayrac
Po-Sen Huang
Robert Stanforth
Alhussein Fawzi
Pushmeet Kohli

title: Are labels required for improving adversarial robustness

authors:
Prasetya Ajie Utama
Nafise Sadat Moosavi
Victor Sanh
Iryna Gurevych

title: Avoiding inference heuristics in few-shot prompt-based finetuning, 2021

authors:
Thomas Wolf
Lysandre Debut
Victor Sanh
Julien Chaumond
Clement Delangue
Anthony Moi
Pierric Cistac
Tim Rault
Rémi Louf
Morgan Funtowicz
Joe Davison
Sam Shleifer
Patrick von Platen
Clara Ma
Yacine Jernite
Julien Plu
Canwen Xu
Teven Le Scao
Sylvain Gugger
Mariama Drame
Quentin Lhoest
Alexander M. Rush

title: Transformers: State-of-the-art natural language processing

venue: In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Pro-

authors:
Mitchell Wortsman
Gabriel Ilharco
Mike Li
Jong Wook Kim
Hannaneh Hajishirzi
Ali Farhadi
Hongseok Namkoong
Ludwig Schmidt

title: Robust fine-tuning of zero-shot models

venue: arXiv preprint arXiv:2109.01903,

authors:
Sang Michael Xie
Ananya Kumar
Robert Jones
Fereshte Khani
Tengyu Ma
Percy Liang

title: InN-out: Pre-training and self-training using auxiliary information for out-of-distribution robustness

venue: In International Conference on Learning Representations (ICLR),

authors:
Elad Ben Zaken
Shauli Ravfogel
Yoav Goldberg

title: Bitfit: Simple parameter-efficient finetuning for transformer-based masked language-models

authors:
Kaiyang Zhou
Jingkang Yang
Chen Change Loy
Ziwei Liu

title: Learning to prompt for visionlanguage models, 2021

title: Human behavior and the principle of least effort

text:
1 INTRODUCTION
When finetuning a pretrained language model for natural language generation tasks like summarization (Narayanan et al., 2018) and table-to-text (Gardent et al., 2017), one is currently faced with a tradeoff. One can achieve strong in-distribution (ID) performance or strong out-of-distribution (OOD) performance—but not both—by choosing one of two families of finetuning approaches. Finetuning all parameters of the pretrained language model achieves strong ID performance (Howard & Ruder, 2018). We call this full finetuning. Alternatively, freezing most of the pretrained parameters during finetuning, in e.g., adapters (Rebuffi et al., 2017; Houlsby et al., 2019a), prefix-tuning (Li & Liang, 2021; Lester et al., 2021) and bitfit (Zaken et al., 2021), achieves stronger OOD performance than full finetuning (Li & Liang, 2021; Lester et al., 2021), but worse ID performance. We call these methods lightweight finetuning.
In this work, we are the first to demonstrate that this tradeoff is not necessary in natural language generation: we can achieve the best of full and lightweight finetuning both ID and OOD. We test this across three
natural language generation tasks in English: summarization (Narayanan et al., 2018), table-to-text (Gardent et al., 2017), and open-domain question answering (Roberts et al., 2020; Lee et al., 2019a), and two lightweight finetuning methods: adapters and prefix-tuning. Achieving the best of both worlds is important since across our 6 settings (2 lightweight finetuning methods and 3 datasets),
full finetuning significantly outperforms lightweight finetuning ID in 5 settings, and lightweight finetuning significantly outperforms full finetuning OOD in 6 settings.
We present two methods for achieving the best of full and lightweight finetuning. We first show that a weighted ensemble of one full finetuning model and one lightweight finetuning model achieves performance comparable to the best of the two, both ID and OOD. To avoid running two models at once, we propose cocktail finetuning: first train a lightweight finetuning model, and then train a full finetuning model on a mixture of student-teacher distillation loss (Bucilă et al., 2006; Hinton et al., 2015) (with the lightweight model as teacher) and cross-entropy loss, on the original ID data.1 The ensemble weight (resp. the cocktail mixture loss weight) is chosen to optimize for ID performance, and we observe that the resultant model achieves comparable OOD performance to lightweight finetuning.
Our work contributes to an ongoing realization in machine learning: robust models that deviate minimally from a pretrained model can be gainfully combined with a model adapted more precisely to a specific dataset. Most saliently, in concurrent work, Wortsman et al. (2021) find that zeroshot image recognition through CLIP (Radford et al., 2021) can be combined with a finetuned network to perform well ID and OOD. Our work demonstrates that similar results hold true in natural language generation as well. We draw from simple, well-known methods: ensembles, student-teacher distillation (Abnar et al., 2020; Hinton et al., 2015)—but our application to achieving the best of ID and OOD performance in natural language generation is novel.
To better understand how cocktail finetuning can transfer the out-of-distribution behavior of the lightweight finetuning model when distilling only on in-distribution data, we theoretically study multiclass logistic regression with a large number of output classes, viewing language generation as a classification problem over a large output space. We study a distribution shift setting where not all classes are seen in the ID training data (as is true in language generation due to Zipf’s law (Zipf, 1949)), and getting good OOD accuracy requires generalizing to unseen outputs. First, we prove that a model trained with gradient descent cannot distinguish the unseen output classes, leading to poor OOD accuracy. Second, we prove that distilling a teacher model with good OOD accuracy on only ID training data can transfer the OOD behavior of the teacher to the student, such that the student and teacher make the same prediction on any new input. As a result, the student also has good OOD accuracy.
In summary, we present two simple methods—ensembling and cocktail finetuning—for achieving strong ID and OOD natural language generation with only the original ID data, achieving the “best of both worlds” of full and lightweight finetuning methods.

heading: 2 RELATED WORK

text:
Lightweight finetuning. Recently, as pretrained models scale up in the number of parameters, lightweight finetuning techniques have been developed to avoid updating all of the parameters. These techniques update a subset of the parameters (Zaken et al., 2021), introduce a small number of new parameters between layers of the model (called “adapters”) (Rebuffi et al., 2017; Houlsby et al., 2019b; Pfeiffer et al., 2020), or learn auxiliary inputs that are prepended to the models inputs (called “prompts”) (Shin et al., 2020; Li & Liang, 2021; Qin & Eisner, 2021; Liu et al., 2021b; Logan IV et al., 2021; and see Liu et al., 2021a for a survey). These techniques empirically have the benefit of being able to achieve strong out-of-distribution extrapolation performance at the cost of worse in-distribution performance compared to full finetuning models both in classification and generation settings (Li & Liang, 2021; Lester et al., 2021; Zhou et al., 2021; Utama et al., 2021). In our work, we aim to overcome this trade-off by proposing techniques for achieving the best of both worlds, with a focus on language generation.
Student-teacher distillation. Buciluǎ et al. (2006) introduce distillation for compressing large ensembles into smaller, faster, more efficient models, and Hinton et al. (2015) generalize their method and call it knowledge distillation. Kasai et al. (2021) distill Transformers into RNNs for linear-time inference. Abnar et al. (2020) show that knowledge distillation can transfer the inductive bias of one architecture to another: for example, distilling an LSTM to a Transformer leads to a Transformer with
1We use the term cocktail as it is a refreshing mixture involving distillation.
representations that are more similar to independently trained LSTMs than independently trained Transformers. The efficiacy of our proposed cocktail finetuning is surprising even in context of these results, as it empirically transfers just the strong OOD performance of the teacher lightweight model to the student, while not transferring the weak ID performance.
Robust finetuning in vision. In concurrent work, Wortsman et al. (2021) propose a similar approach to ours, but focus on image classification: they combine a full finetuning model with a lightweight finetuning (linear probe) model by ensembling the two models in weight space, and show that their approach maintains both good ID and OOD performance. It is yet surprising that in the structured prediction setting of natural language generation, we further show that we can achieve high OOD performance even when weight-space ensembling is impossible (for instance, since there are weights not held in common between the full finetuning model and an optimized prompt). Moreover, we further this line of work by providing a theoretical study into how distillation can transfer OOD behavior while only using ID data.
Self-training for robustness. In adversarial robustness, standard models achieve good accuracy on clean examples but are not robust to adversarial examples, while adversarially trained models are more robust but are less accurate on clean examples. Raghunathan et al. (2020) show that the benefits of these two models can be combined using robust self-training (RST) (Carmon et al., 2019; Uesato et al., 2019; Najafi et al., 2019) on unlabeled data to improve clean and robust accuracy. Similarly, Khani & Liang (2021) show that RST can be used to combine models with spurious features removed (which are more robust) and without removing spurious features (which are more accurate). In general OOD robustness, In-N-Out (Xie et al., 2021) uses self-training to transfer the robustness benefits of a model pre-trained with auxiliary self-supervision to a model trained with extra auxiliary features. Berthelot et al. (2021) use unlabeled data and a consistency regularization loss for domain adaptation. While many of these methods combine the benefits of two models or use a self-training/distillation loss to improve robustness, cocktail finetuning differs in that it performs distillation only on ID training data without any additional ID or OOD unlabeled examples.

heading: 3 PROBLEM STATEMENT

heading: 3.1 SETUP

text:
We consider a conditional generation task from an input space X (e.g., a news article) to an output space Y (e.g., a summary of the article), where both spaces consist of sequences of tokens from a shared vocabulary V .
Data. Let Pid and Pood denote the distribution of (x, y) pairs in-distribution and out-of-distribution, respectively, where x ∈ X and y ∈ Y . In summarization, Pid may consist of articles about world and business news while Pood consists of health and technology news. The training data Dtrain and validation data Dval consist of n and nval in-distribution (x, y) pairs respectively, drawn from Pid. Model and metrics. Let qγ(y | x) = ∏len(y) j=1 qγ(yj | x, y<j) be a pretrained neural autoregressive sequence model with parameters γ. Finetuning (full or lightweight) produces a finetuned model pθ(y | x) from qγ , where θ denotes the set of trainable parameters. In full finetuning, θ is the same set of parameters as γ; in lightweight finetuning, θ often refers to the parameters of inserted modules (e.g., adapters modules or prefixes). The final predictor is fθ(x) which approximates argmaxy pθ(y | x) through, e.g., beam search. The goal is to learn parameters θ such that the predictor fθ obtains good ID and OOD performance:
Sid(fθ) = Ex,y∼Pid [score(y, fθ(x))] (1) Sood(fθ) = Ex,y∼Pood [score(y, fθ(x))] (2)
where score may be defined differently for each task (e.g., exact match score or BLEU (Papineni et al., 2002b)).

heading: 3.2 BASELINE FINE-TUNING METHODS

text:
Full finetuning has been the defacto approach to adapt pretrained language models for downstream generation tasks (Lewis et al., 2020; Kale & Rastogi, 2020). Under this framework, the model pθ is initialized with its pretrained parameters γ, and we optimize the following log-likelihood objective.
L(θ) = E(x,y)∼Dtrain [− log pθ(y | x)] (3)
where the set of trainable parameters θ refer to the set of pretrained parameters γ.
Lightweight finetuning is another family of adaptation approaches, which freezes most of the pretrained parameters and augments the model with small trainable modules. Lightweight finetuning shares the same training objective as full finetuning, but the trainable parameters θ are different. Within the lightweight finetuning family, we focus on adapter-tuning (Houlsby et al., 2019a) and prefix-tuning (Li & Liang, 2021). Both approaches keep the pretrained parameters intact. Adaptertuning inserts task-specific MLP layers between each layer of the pretrained language models, and the trainable parameter θ consists of these inserted MLP parameters. Prefix-tuning prepends a sequence of trainable, task-specific prefix vectors to the input, and θ consists of these prefix parameters.2

heading: 4 EXPERIMENT SETUP

heading: 4.1 DATASETS AND METRICS

text:
Table-to-text. We evaluate on the WebNLG dataset (Gardent et al., 2017), where the input x is a sequence of (entity, predicate, entity) triples and the output y is a natural language description that covers the input information. Each example in this dataset is labeled with one of 14 topic (e.g., airports, comic characters). We select 9 topics as in-distribution and the remaining 5 topics as out-of-distribution. We evaluate the performance using the official evaluation scripts and report the BLEU scores (Papineni et al., 2002a).
Summarization. We use the XSUM (Narayan et al., 2018) dataset, which is an abstractive summarization dataset on BBC news articles. Each example is labeled with a news category. We select world, UK and business categories to be in-distribution and let the remaining news categories such as health and technology to be out-of-distribution. We report the ROUGE-2 scores (Lin, 2004).
Open-domain, closed-book QA (Open-QA). We use the open-domain variant (Lee et al., 2019b) of the Natural Questions dataset (Kwiatkowski et al., 2019) as the in-distribution data and the Web Questions dataset (Berant et al., 2013) as the out-of-distribution data. The distribution shift is induced by different data collection process and their year of collection (roughly 2018 versus 2013). We
2For specifics of these methods, see their respective papers.
report the exact match accuracy. Note that this task, unlike most QA (e.g., SQuAD (Rajpurkar et al., 2016)), does not provide systems with a text from which to extract answers. Roberts et al. (2020) demonstrated that one could finetune an LM to generate answers from knowledge stored during pretraining.

heading: 4.2 MODEL ARCHITECTURES AND HYPERPARAMETERS

text:
We use GPT-2-medium (Radford et al., 2019) as the pretrained model for table-to-text and opendomain QA tasks. GPT-2 is an autoregressive language model; therefore, we parametrize pθ(y | x) by concatenating x, a separator token, and y (e.g., [x; SEP; y]). We define the loss function to sum over the token losses that correspond to y.
We use BART-large (Lewis et al., 2020) as the pretrained model for the summarization task. BART uses an encoder-decoder architecture and parametrizes pθ(y | x) by first feeding x to the encoder, and decoding y conditioned on the encoder representation.
For open-domain closed-book QA, this is to our knowledge the first time GPT-2 has been finetuned for the task. We used separate development OOD data to test a handful (<10) of hyperparameter settings to determine whether we’d observe a similar ID/OOD tradeoff to the other tasks. We note this because, in general, one cannot do hyperparameter optimization on OOD data; we did this to provide another testbed for our methods.
Our implementation is based on the Hugging Face Transformers package (Wolf et al., 2020). We use most of the hyperparameters suggested by the Hugging Face default setup. The hyperparameters we tune include the number of epochs, batch size, learning rate, and prefix length (for prefix-tuning), all reported in Table 3 in the appendix.

heading: 5 RESULTS OF FULL AND LIGHTWEIGHT FINETUNING

text:
There is a trade-off in ID/OOD performance when choosing between full finetuning and lightweight finetuning: full finetuning consistently outperforms in-distribution, whereas lightweight finetuning consistently outperforms out-of-distribution (Li & Liang, 2021; Lester et al., 2021). To quantify this trade-off, we examine the generation performance of both methods on ID and OOD data.
As shown in Table 2, when evaluating in-distribution, full finetuning achieves significantly stronger results than lightweight finetuning in five out of six settings (across 3 datasets and two lightweight finetuning methods. Only prefix-tuning on WebNLG achieves comparable accuracy to full finetuning. When evaluating out-of-distribution, lightweight finetuning achieves significantly stronger results than full finetuning in five of six settings. Only adapter-tuning on XSUM fails to outperform full finetuning.3

heading: 6 COMBINING MODELS WITH A WEIGHTED ENSEMBLE

text:
As we observe in Section 5, full finetuning models attain better ID performance but worse OOD performance, whereas lightweight finetuning models attain better OOD performance at the cost of good ID performance. Is it possible to achieve both good ID and OOD performance simultaneously? In this section, we show that a simple ensemble of the two methods accomplishes this.
First, we train a lightweight finetuning model p(`) and a full finetuning model p(f) on Dtrain. At decoding time, we combine the prediction of the p(`) and p(f) by interpolating them at each token.
p(Ens)(yt | x, y<t) = λ · p(`)(yt | x, y<t) + (1− λ) · p(f)(yt | x, y<t) (4)
At each time step t, both models output a next token distribution conditioned on the common history y<t, and the two distributions are interpolated with a mixture weight λ. The output is generated by beam searching the ensemble distribution p(Ens)(yt | x, y<t).
3Though it is worth noting again that for open-domain QA, we performed a small amount of manual, exploratory hyperparameter optimization to optimize the ID of full finetuning and the OOD of lightweight finetuning.
6.1 CHOOSING THE MIXTURE WEIGHT λ
We start with a set of {λi}Ki=1 where λi ∈ [0, 1]. For example we use λ ∈ {0., 0.1, 0.25, 0.5, 0.75, 0.9, 1.0} in this experiment. We choose the mixture weight λi that attains the best score on Dval.
This method of choosing λ is useful because (1) it crucially does not use OOD data, which is typically unaccessible at model selection time, and (2) it guarantees that ID performance is not sacrificed, because we can match ID performance by setting λ = 0.

heading: 6.2 RESULTS

text: As shown in Table 2, the ensemble matches and slightly outperforms both the ID performance skyline (achieved by full finetuning) and the OOD performance skyline (achieved by lightweight finetuning). Across all datasets, ensembling achieves significantly improved performance over full finetuning OOD, and lightweight finetuning ID. Further, in eleven out of twelve settings (across 3 datasets, 2 methods, and both ID and OOD), ensembling does not perform significantly worse than the better of full and lightweight finetuning.

heading: 6.3 RUNTIME AND MEMORY

text: Despite its ability to combine the best of both worlds, ensembling is expensive at inference time because two models are run simultaneously, doubling the amount of GPU memory required. Furthermore, two forward passes are required, doubling the runtime for decoding.4 For deployment efficiency, it’s better to have a single model at inference time.

heading: 7 COCKTAIL FINETUNING: COMBINING WITH DISTILLATION

text: The goal of cocktail finetuning is to approximate the ensemble in Section 6 with a single model.5

heading: 7.1 METHOD

text:
Train a lightweight finetuning model p(`). Fix mixture weight λ ∈ [0, 1]. Train all parameters of the pretrained model qγ(y | x) on the following objective. First, define a token-level student-teacher
4The two forward passes are parallelizable, but at the cost of another 2x GPU memory. 5One might wish to approximate the two with a single lightweight model, a direction we explored but eventually discarded due to the difficulty of approximating a full finetuning model with a lightweight finetuning model. Anecdotally, this is evidence towards the speculation of Lester et al. (2021), that lightweight finetuning is limited in its ability to alter a pretrained LM.
distillation loss at token j: distillp(`)(x, y, j) = KL ( p(`)(yj | x, y<j)‖pθ(yj | x, y<j)) ) (5)
The objective is the sum over the length of y of the λ-weighted mixture of distillation and loglikelihood loss:
Lλ(θ) = E(x,y)∼Dtrain [ len(y)∑ j=1 λdistillp(`)(x, y, j)− (1− λ) log pθ(yj | x, y<j) ]
(6)
Let fθ,λ be the predictor resulting from optimizing on Lλ(θ) and estimate the ID score Sid(fθ,λ) on Dval. Perform this across a set of mixture weights {λi}i. The cocktail finetuning model is the predictor fθ,λi that maximizes the in-distribution score Sid(fθ, λi) across the mixture weights λi. 6

heading: 7.2 RESULTS

text: As shown in Table 2, cocktail finetuning matches and slightly outperforms the ID performance skyline (achieved by full finetuning) and achieves comparable to or slightly worse than the OOD performance skyline (achieved by lightweight finetuning). Across all datasets, cocktail finetuning achieves significantly improved performance over full finetuning OOD, and over lightweight finetuning ID. Further, in seven out of twelve settings (across 3 datasets, 2 methods, and both ID and OOD), cocktail finetuning does not perform significantly worse than the better of full and lightweight finetuning.

heading: 8 ANALYSIS OF THE WEIGHTING OF FULL AND LIGHTWEIGHT FINETUNING

text:
Trade-off curves. Our ensembling and cocktail finetuning methods are useful in practice because one can choose the mixture coefficient λ without OOD data. Now, we explore how ID and OOD performance vary as a function of λ. Figure 3 shows results across five random seeds of ensembling and cocktail finetuning on our three datasets, each dot colored according to its λ value. In each plot, we see a smooth tradeoff curve, in which a range of λ improve on the OOD, ID, or both of each constituent model. Interestingly, in two of three tasks, there is effectively only one pareto-optimal λ. This means the λ is best for both for ID and OOD
performance, which is why optimizing for ID performance also leads to good OOD performance.
6For cocktail finetuning, we use the hyperparameters reported for full finetuning; no hyperparameter optimization was performed.
Ablations. One hypothesis to explain the strong performance of ensembles is the known benefits of ensembling two models trained with different random seeds. Likewise for cocktail finetuning, one could point to the effectiveness of born-again networks (Furlanello et al., 2018): distilling a teacher model to an identically parameterized student with the same model architecture. As an ablation to test this hypothesis, we distill from a full finetuning model into another full finetuning model, again using λ to mix the distillation loss with the label loss. The results, in Figure 4, show that this distillation does not improve over a single full finetuning model, for any of the λ.

heading: 9 WHY DOES DISTILLATION ON ID DATA IMPROVE OOD?

text:
Prior works have used unlabeled OOD data and self-training (similar to distillation) to transfer good OOD performance from another model (Raghunathan et al., 2020; Xie et al., 2021). In contrast, cocktail finetuning performs distillation only on ID data. In this section, we provide an explanation for why this can occur in multiclass logistic regression with a large number of classes.
Prediction problem and distribution shift. Motivated by the large output space in language generation, we consider a prediction problem from inputs X ⊆ Rd to labels Y = {1, . . . , k} where k is large. The labeled training data is {(xi, yi)}ni=1. For convenience, we also let yi be the one-hot vector of length k denoting the correct class. The labels in the training data do not cover all the possible labels, and we let the subset of seen labels be S. In this problem, the class marginals of the ID and OOD distributions Pid and Pood are different. To have good OOD accuracy, the model must generalize to unseen classes.
Model and losses. We consider training a multi-class logistic regression model using full batch gradient descent where the learned parameters are θ ∈ Rk×d. In full fine-tuning, we run gradient descent on the negative log-likelihood.
In cocktail finetuning, we also add a distillation loss using soft labels (a probability vector) from the prefix-tuned model, which we denote px
θ̃ = pθ̃(· | x) ∈ Rk with parameters θ̃, given an input x.
Similarly, the predicted class probability vector of the cocktail model given an input x is denoted
pxθ = pθ(· | x) ∈ Rk. In this setting, the cocktail loss (Equation 6) with parameter λ ∈ [0, 1] is
`ct(θ) = n∑ i=1 ((1− λ)yi + λpxiθ̃ ) > log pxiθ . (7)
The cocktail loss encourages the model to learn to output a mixture of the hard label and the soft label on the training data. The final score for the classification task is computed with respect to the accuracy metric, where score(y, ŷ) = 1− `0−1(y, ŷ) and `0−1 is the 0-1 loss.
Poor OOD performance with full finetuning. Our first result shows that when training a logistic regression model using gradient descent (initialized at 0) on the only labeled training set (as in full finetuning), the learned parameters θj for an unseen label j /∈ S are all the same.
Proposition 1. Let θ(t) be the parameters at iteration t of gradient descent and let θ(0) = 0. At any iteration t and for any unseen label index j /∈ S, θj = gt where gt is the same for all j /∈ S.
As a result, the model cannot distinguish between any of the unseen classes, resulting in poor OOD performance. Suppose that the model outputs a fixed, arbitrary unseen class j ∈ S when the probabilities of the unseen classes are the largest. Then, assuming that the OOD distribution Pood is such that the probability of drawing an example from an unseen class (excluding j) Pood(y /∈ S∪{j}) is at least c ∈ [0, 1], then the OOD accuracy is bounded as
Sood ≤ 1− Pood(y /∈ S ∪ {j}) ≤ 1− c (8)
This OOD accuracy can be low when c is large, which can occur with very large output spaces.
Disillation transfers OOD behavior through soft logits. We investigate how distillation on ID training data can affect OOD behavior. In the following, let X ∈ Rn×d be the training data matrix. Proposition 2. Assume the data matrixX has full column rank. Let θ be the minimizer of the cocktail loss when λ = 1, subject to the constraint that 1> exp(θX) = α1> exp(θ̃X) for some constant α > 0, such that the model learns the same normalizing constants (up to scaling) as the prefix-tuned model on the training data only. Then
pθ(· | x) = pθ̃(· | x) (9)
for any new input x, such that the cocktail model has the same OOD accuracy as the prefix model.
Therefore, if the prefix-tuning model has high OOD accuracy, then Sood is high. The proof proceeds by showing that the extra information in the soft logits can help the model recover the parameters of the prefix-tuned model (up to scaling), which in turn inherits its OOD behavior. Thus, in generation problems where the output space is large, distillation on soft logits can transfer information about OOD behavior even through ID data.

heading: 10 CONCLUSION

text:
Considered together, our work and the work of Wortsman et al. (2021) provide strong cross-modal evidence that simple ensembles of models that have minimally deviated from pretraining (for them, zero-shot; for us, lightweight finetuning models) with full finetuning models provides strong ID and OOD accuracy. The reasons for this “best of both worlds” behavior are still poorly understood, motivating future work. Since ensemble models effectively weight by model confidence (though modulated by λ), one explanation may be that lightweight finetuning models are less confident than full finetuning models in-distribution, yet more confident out-of-distribution.
In summary, we provide simple, effective methods for achieving the best of lightweight and full finetuning—strong OOD and ID performance—in language generation. Our cocktail finetuning provides system builders with the opportunity to optimize for in-distribution performance (in choosing lambda), often beating full finetuning, while resting assured that the out-of-distribution performance (i.e., what they observe with real users) will likely be higher than had they performed full finetuning.
ETHICS STATEMENT
The methods presented in this work have the potential to improve a wide range of natural language generation tasks. Natural language generation has known dual-use issues, like scaling misinformation generation and generating naturalistic text in scams. The WebNLG XSUM, Natural Questions, and WebQuestions datasets we evaluate on contain personal information (though public and published) information about real people and organizations.
REPRODUCIBILITY STATEMENT
All code (including experiment scripts and plot creation scripts) will be released upon publication. All datasets used in this paper are publicly available and will be linked in our released code. All reported numbers represent the average of results from 5 random seeds.

heading: A HYPERPARAMETERS AND DATASET DETAILS

text:
A.1 HYPERPARAMETERS
Relevant hyperparameters beyond the defaults of Huggingface Transformers Wolf et al. (2020) are reported in Table 3.
A.2 DATASET DETAILS
Table 5 shows the number of examples in each dataset split.
For the WebNLG dataset, the ID topics include {Astronaut, University, Monument, Building, ComicsCharacter, Food, Airport, SportsTeam, City, and WrittenWork} and the OOD topics include {Athlete, Artist, MeanOfTransportation, CelestialBody, Politician}. For the XSUM dataset, the ID topics include {uk, business, world } news and the OOD topics includes { entertainment, election, technology, science, health, education, explainers, live, ukpolitics, world-us-canada, world-europe, uk-england, uk-scotland, world-asia, science-environment, uk-scotland-scotland-politics, disability}.

heading: B ADDITIONAL RESULTS

text: Due to space constraints, we did not include in Figure 3 the tradeoff plots using adapters lightweight finetuning method (instead reporting only those for prefix-tuning.) The plots for adapters are provided in Figure 5.

heading: C ANALYSIS

text:
We restate and expand upon some of the setup for the analysis here.
Prediction problem. We consider a prediction problem from inputs X ⊆ Rd to labels Y = {1, . . . , k} where k is large. The labeled training data is {(xi, yi)}ni=1. For convenience, we also let yi be the one-hot vector of length k denoting the correct class, and will make this clear. The labels in the training data do not cover all the possible labels, and we let the subset of seen labels be S.
Model. We consider training a multi-class logistic regression model using full batch gradient descent. The learned parameters are θ ∈ Rk×d and the model prediction given input x is
fθ(x) = argmax j pθ(x)j = argmax j θ>j x∑ j′ θ > j′x
(10)
where θj is the j-th row of θ. In standard fine-tuning, we run gradient descent on the negative log-likelihood
`(A) = − n∑ i=1 y>i log p xi θ (11)
where yi ∈ Rk is a one-hot vector of the class label.
Distillation loss. In the Cocktail model, we also add a distillation loss using soft labels from the prefix-tuned model, which we denote pθ̃ with parameters θ̃. The Cocktail loss with parameter λ ∈ [0, 1] is
`ct(A) = n∑ i=1 ( (1− λ)y>i log p xi θ + λ(p xi θ̃ )> log pxiθ ) (12)
= n∑ i=1 ((1− λ)yi + λpxiθ̃ ) > log pxiθ (13)
The Cocktail loss encourages the model to learn to output a mixture of the hard label and the soft label on the training data.
C.1 PROOF OF PROPOSITION 1
Proof. By induction on t. By zero initialization, the statement holds with g0 = 0.
Suppose for iterations up to t, we have θ(t)j = gt for all j /∈ S . We focus on some j /∈ S WLOG. The gradient for θ(t)j is
∇ θ (t) j `(θ(t)) = ∇ θ (t) j
( −
n∑ i=1 y>i log p xi θ
) (14)
= n∑ i=1 −yijxi + exp((θ (t) j ) >xi)∑ j′ exp((θ (t) j′ ) >xi) xi (15)
= n∑ i=1 exp((θ (t) j ) >xi)∑ j′ exp((θ (t) j′ ) >xi) xi (16)
= n∑ i=1 exp(g>t xi)∑ j′ exp((θ (t) j′ ) >xi) xi (17)
since yij = 0 for all i as j is unseen. This gradient does not vary with j, which shows the result.
C.2 PROOF OF PROPOSITION 2
Proof. Let Yθ ∈ Rn×k denote the predicted class probabilities of the model on the training data, and Yθ̃ be the class probabilities of the prefix-tuned model on training data. With λ = 1, the gradient of the loss is
∇θ`ct(θ) = n∑ i=1 −pxi θ̃ x>i + p xi θ x > i (18)
= −Yθ̃X + YθX (19) Since θ is the minimizer, this gradient is 0 and thus
Yθ̃X = YθX =⇒ Yθ̃ = Yθ (20) for X with full column rank (right-invertible). We let the right-inverse of X be X†. This implies that for all i,
exp(θ̃xi)∑ j exp(θ̃ > j xi) = exp(θxi)∑ j exp(θ > j xi) . (21)
In matrix form, this is equivalent to
exp(θ̃X) · diag(1> exp(θ̃X))−1 = exp(θX) · diag(1> exp(θX))−1 (22) =⇒ θ̃X − 1(log 1> exp(θ̃X))> = θX − 1(log 1> exp(θX))> (23)
=⇒ θ = θ̃ + 1 [ log 1> exp(θX) log 1> exp(θ̃X) ] X†. (24)
On a new input x,
θx = θ̃x+ 1 [ log 1> exp(θX) log 1> exp(θ̃X) ] X†x (25)
= θ̃x+ (log(α) + 1>X†x)1 (26)
by using the assumption on normalizing constants. Therefore the probability of the j-th class under the model is
(pxθ )j = exp(θ>j x)∑ j′ exp(θ > j′x)
(27)
= exp(θ̃>j x) exp(log(α) + 1 >X†x)∑ j′ exp(θ̃ > j′x) exp(log(α) + 1 >X†x) (28) = exp(θ̃>j x)∑ j′ exp(θ̃ > j′x) = (px θ̃ )j (29)