summaries: Shapeshifter reparameterizes dense kernels into a product of two lower-rank matrices. Shapeshifter introduces reshapes and transposes to make this reparameterization more efficient. The paper provides a theoretical justification for its expressivity and strong experimental results in machine translation.

summaries: This paper proposes a method to reduce the number of parameters by weight factorization. Unlike traditional low-rank factorization (n x r and r x m), this work decomposes the matrix into two sqrt(mn) x r matrices. The algorithm includes the reordering of elements to recover the desired shape. The authors find connections to Kronecker products and theoretically support the expressiveness of the proposed method. Experiments on NMT tasks show that Shapeshifter can achieve better parameter efficiency (fewer parameters with comparable performance). The results are compared to two similar works: PHM and DeLighT.

summaries: This paper presents an approach called ShapeShifter to compress large language models. The paper follows the approach of reducing number of parameters by using factorized matrix representations. The primary idea is to represent the matrix as a sum of Kronecker products, instead of standard low-rank factorizations. The paper theoretically shows that for non-square matrices, the proposed approach allows for increased reduction in parameters for the same decomposition rank. For square matrices, it allows for increased expressiveness of the decomposition for the same rank. Experimental results on machine translation tasks show that the proposed approach performs much better than the state-of-the-art model compression approaches and is just slightly worse than the large models.

summaries: A sizable effort has been made in recent history to obtain the accuracy of very large language models at a fraction of the parameter of the large models. Broadly these methods fall into various categories like model compression, factorized representation, distillation etc. This paper underlines another such attempt at obtaining a factorized representation of matrices to reduce the parameter space. The authors justify the use of Kronecker products to factorize parameter space, by showing that stacked layers of low rank matrices increase expressiveness. They also show the effectiveness of their proposed method by achieving a 4 to 14-fold reduction in parameters without sacrificing accuracy.

summaries: This paper proposes an algorithm they call ‘Shapeshifter’ to reduce the number of parameters needed for Transformer-based models while preserving expressiveness and performance. The proposed algorithm is based on the idea of low-rank factorizations of weight matrices, but uses a different factorization method involving sums of Kronecker products that leads to greater expressiveness. They show theoretical analysis to prove their factorization can accurately represent any $n \times m$ matrix or stack of matrices (in an idealized case), and experimental results using transformer-based models for translation on English-German, English-French and English-Romanian. They demonstrate superior or comparable performance with vastly fewer parameters compared to some other model compression techniques.

abstractText: Language models employ a very large number of trainable parameters. Despite being highly overparameterized, these networks often achieve good out-of-sample test performance on the original task and easily fine-tune to related tasks. Recent observations involving, for example, intrinsic dimension of the objective landscape and the lottery ticket hypothesis, indicate that often training actively involves only a small fraction of the parameter space. Thus, a question remains how large a parameter space needs to be in the first place — the evidence from recent work on model compression, parameter sharing, factorized representations, and knowledge distillation increasingly shows that models can be made much smaller and still perform well. Here, we focus on factorized representations of matrices that underpin dense, embedding, and self-attention layers. We use low-rank factorized representation of a reshaped and rearranged original matrix to achieve space efficient and expressive linear layers. We prove that stacking such low-rank layers increases their expressiveness, providing theoretical understanding for their effectiveness in deep networks. In Transformer models, our approach leads to more than tenfold reduction in the number of total trainable parameters, including embedding, attention, and feed-forward layers, with little degradation in on-task performance. The approach operates out-of-the-box, replacing each parameter matrix with its compact equivalent while maintaining the architecture of the network.

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

title: Attention is all you need

venue: In Proceedings of the 31st International Conference on Neural Information Processing Systems,

authors:
Jacob Devlin
Ming-Wei Chang
Kenton Lee
Kristina Toutanova

title: BERT: pre-training of deep bidirectional transformers for language understanding

venue: Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies,

authors:
Tom Brown
Benjamin Mann
Nick Ryder
Melanie Subbiah
Jared D Kaplan
Prafulla Dhariwal
Arvind Neelakantan
Pranav Shyam
Girish Sastry
Amanda Askell
Sandhini Agarwal
Ariel Herbert-Voss
Gretchen Krueger
Tom Henighan
Rewon Child
Aditya Ramesh
Daniel Ziegler
Jeffrey Wu
Clemens Winter
Chris Hesse
Mark Chen
Eric Sigler
Mateusz Litwin
Scott Gray
Benjamin Chess
Jack Clark
Christopher Berner
Sam McCandlish
Alec Radford
Ilya Sutskever
Dario Amodei

title: Language models are few-shot learners

venue: Advances in Neural Information Processing Systems,

authors:
William Fedus
Barret Zoph
Noam Shazeer

title: Switch Transformers: Scaling to trillion parameter models with simple and efficient sparsity

venue: arXiv preprint arXiv:2101.03961,

authors:
Chunyuan Li
Heerad Farkhoor
Rosanne Liu
Jason Yosinski

title: Measuring the intrinsic dimension of objective landscapes

venue: In International Conference on Learning Representations,

authors:
Armen Aghajanyan
Sonal Gupta
Luke Zettlemoyer

title: Intrinsic dimensionality explains the effectiveness of language model fine-tuning

venue: In 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),

authors:
Jonathan Frankle
Michael Carbin

title: The lottery ticket hypothesis: Finding sparse, trainable neural networks

authors:
Jonathan Frankle
Gintare Karolina Dziugaite
Daniel Roy
Michael Carbin

title: Linear mode connectivity and the lottery ticket hypothesis

venue: In International Conference on Machine Learning,

authors:
Tianlong Chen
Jonathan Frankle
Shiyu Chang
Sijia Liu
Yang Zhang
Zhangyang Wang
Michael Carbin

title: The lottery ticket hypothesis for pre-trained BERT networks

venue: In Advances in Neural Information Processing Systems,

authors:
Sai Prasanna
Anna Rogers
Anna Rumshisky

title: When BERT plays the lottery, all tickets are winning

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

authors:
Victor Sanh
Lysandre Debut
Julien Chaumond
Thomas Wolf

title: DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter

authors:
Bei Li
Ziyang Wang
Hui Liu
Quan Du
Tong Xiao
Chunliang Zhang
Jingbo Zhu

title: Learning light-weight translation models from deep transformer

venue: Proceedings of the AAAI Conference on Artificial Intelligence,

authors:
Fusheng Wang
Jianhao Yan
Fandong Meng
Jie Zhou

title: Selective knowledge distillation for neural machine translation

authors:
Sachin Mehta
Marjan Ghazvininejad
Srinivasan Iyer
Luke Zettlemoyer
Hannaneh Hajishirzi

title: DeLighT: Very deep and light-weight transformer

venue: In International Conference on Learning Representations ICLR’2021,

authors:
Zhenzhong Lan
Mingda Chen
Sebastian Goodman
Kevin Gimpel
Piyush Sharma
Radu Soricut

title: ALBERT: A lite BERT for self-supervised learning of language representations

venue: In International Conference on Learning Representations

authors:
Minh-Thang Luong
Hieu Pham
Christopher D Manning

title: Effective approaches to attentionbased neural machine translation

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

authors:
Aliakbar Panahi
Seyran Saeedi
Tom Arodz

title: word2ket: Space-efficient word embeddings inspired by quantum entanglement

authors:
Aston Zhang
Yi Tay
SHUAI Zhang
Alvin Chan
Anh Tuan Luu
Siu Hui
Jie Fu

title: Beyond fully-connected layers with quaternions: Parameterization of hypercomplex multiplications with 1/n parameters

authors:
Yi Tay
Aston Zhang
Anh Tuan Luu
Jinfeng Rao
Shuai Zhang
Shuohang Wang
Jie Fu
Siu Cheung Hui

title: Lightweight and efficient neural natural language processing with quaternion networks

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

authors:
Shuchang Zhou
Jia-Nan Wu
Yuxin Wu
Xinyu Zhou

title: Exploiting local structures with the Kronecker layer in convolutional networks

venue: arXiv preprint arXiv:1512.09194,

authors:
Cijo Jose
Moustapha Cissé
Francois Fleuret

title: Kronecker recurrent units

title: Tensor-train decomposition

venue: SIAM Journal on Scientific Computing,

authors:
Emily Denton
Wojciech Zaremba
Joan Bruna
Yann LeCun
Rob Fergus

title: Exploiting linear structure within convolutional networks for efficient evaluation

authors:
Alexander Novikov
Dmitry Podoprikhin
Anton Osokin
Dmitry Vetrov

title: Tensorizing neural networks

authors:
Charles C Onu
Jacob E Miller
Doina Precup

title: A fully tensorized recurrent neural network

venue: arXiv preprint arXiv:2010.04196,

authors:
Oleksii Hrinchuk
Valentin Khrulkov
Leyla Mirvakhabova
Elena Orlova
Ivan Oseledets

title: Tensorized embedding layers

venue: In Findings of the Association for Computational Linguistics: EMNLP 2020,

authors:
Xindian Ma
Peng Zhang
Shuai Zhang
Nan Duan
Yuexian Hou
Ming Zhou
Dawei Song

title: A tensorized transformer for language modeling

authors:
Thomas Frerix
Joan Bruna

title: Approximating orthogonal matrices with effective Givens factorization

authors:
Harold V Henderson
Shayle R Searle

title: The vec-permutation matrix, the vec operator and Kronecker products: A review

venue: Linear and multilinear algebra,

title: The ubiquitous Kronecker product

venue: Journal of Computational and Applied Mathematics,

authors:
Charles F Van Loan
Nikos Pitsianis

title: Approximation with Kronecker products

authors:
Julie Kamm
James G Nagy

title: Kronecker product and SVD approximations in image restoration

venue: Linear Algebra and its Applications,

title: Kronecker-product approximations for some function-related matrices

title: Computation of plain unitary rotations transforming a general matrix to triangular form

venue: Journal of the Society for Industrial and Applied Mathematics,

authors:
Colin Raffel
Noam Shazeer
Adam Roberts
Katherine Lee
Sharan Narang
Michael Matena
Yanqi Zhou
Wei Li
Peter J. Liu

title: Exploring the limits of transfer learning with a unified text-to-text transformer

venue: Journal of Machine Learning Research,

authors:
Ond Onrej Bojar
Christian Federmann
Mark Fishel
Yvette Graham
Barry Haddow
Matthias Huck
Philipp Koehn
Christof Monz

title: Findings of the 2018 Conference on Machine Translation (WMT18)

venue: In Proceedings of the Third Conference on Machine Translation,

authors:
Mauro Cettolo
J Niehues
S Stüker
Luisa Bentivogli
Marcello Federico

title: Report on the 11th IWSLT evaluation campaign, IWSLT

venue: In IWSLT-International Workshop on Spoken Language Processing,

authors:
Rico Sennrich
Barry Haddow
Alexandra Birch

title: Neural machine translation of rare words with subword units

venue: In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers),

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 of the Association for Computational Linguistics,

authors:
Thomas Wolf
Lysandre Debut
Victor Sanh
Julien Chaumond
Clement Delangue
Anthony Moi
Pierric Cistac
Tim Rault
Rémi Louf
Morgan Funtowicz

title: HuggingFace’s Transformers: State-of-the-art natural language processing

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:
Yang You
Jing Li
Sashank Reddi
Jonathan Hseu
Sanjiv Kumar
Srinadh Bhojanapalli
Xiaodan Song
James Demmel
Kurt Keutzer
Cho-Jui Hsieh

title: Large batch optimization for deep learning: Training BERT in 76 minutes

venue: In International Conference on Learning Representations ICLR’2020,

heading: 1 Introduction

text:
Natural language models involve large number of parameters. A single encoder-decoder Transformer [1] in its base variant has about 44 million parameters, not counting the word embedding matrix, which adds another 10 million or more, depending on the chosen vocabulary or tokenization scheme. Base variant of encoder-only BERT [2], including the embedding, has about 108 million parameters. GPT-3 [3] has about 175 billion parameters, and the largest of the Switch Transformer [4] models has 1.5 trillion. This explosion in the model size has led to increased interest in approaches for reducing the number of parameters in the model.
Models with high-dimensional parameter space have much lower intrinsic dimension [5], that is, training trajectory can be successfully restricted to a random, smaller-dimensional subspace, even ∗Work performed while at Virginia Commonwealth University. †Corresponding author.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
though training a small-parameter architecture is often less successful. These observations have been recently extended to the fine-tuning trajectories of language models [6]. Lottery ticket hypothesis [7, 8], recently demonstrated to hold also for language model fine-tuning [9, 10], shows that smaller subnetworks can be selected from a large model, re-trained in isolation, and perform as well as the large model; what those subnetworks are is not known a prior, in absence of the trained large model, though. Some approaches for reducing model size build on this observation to train a smaller model based on an existing large model, for example a 66 million parameter DistillBERT [11] student has been distilled from 108 million parameter BERT-base [2] teacher with little loss in quality. Alternatively, distillation has been used, for example in the context of machine translation, to increase the quality of the models while maintaining their size [12, 13]. Other approaches train a reduced-parameter model de novo, without relying on an already trained large model. For example, DeLighT [14] uses an alternative parameterization of the multi-headed self-attention based on group linear transform to reduce a 62 million parameter Transformer to 22 million.
One simple way to reduce model size involves factorized matrix representations. ALBERT [15] employs a rank r decomposition of a d×nvocab embedding matrix storing d-dimensional embedding vectors for each of the nvocab tokens by using a stack of two linear layers, d× r on top of r× nvocab. Similar low-rank decomposition is also used implicitly in the multi-headed self-attention in the generic Transformer [1] with hidden dimension d and nheads self-attention heads. In each Transformer head, a r = d/nheads-rank factorized representation involving d × d/nheads key (K) and query (Q) matrices are used, with the pairwise self-attention scores for sequence x calculated using xTKTQx, instead of xTWx involving a full general attention d× d trainable matrix W as originally considered in trainable-attention encoder-decoder LSTM models [16]. In both cases, the models are trained from a random initialization of the factorized parameter space, instead of attempting to find the lowest-error factorized representation of an already trained original model.
We explore here a Transformer model that uses an alternative way to decompose a matrix into two smaller matrices. Instead of standard low-rank factorization as above, it involves reshaping and reordering matrix dimensions prior to the decomposition, and is equivalent to a sum of Kronecker products with an efficient implementation. The compact model based on the smaller matrices is trained de novo, and matrices in the compact model are not aimed to be approximations of the matrices in the original model. For non-square model matrices, the approach allows for increased reduction in parameters for the same decomposition rank. For square matrices, the benefits come from increased expressiveness of the decomposition for the same rank, allowing for reducing the rank needed to preserve model accuracy, and thus reducing the model size. Our main contribution is proving that stacking multiple linear layers decomposed this way increases the expressiveness of the network, unlike stacking multiple low-rank layers factorized in the standard way, which can only map into a subspace of dimensionality equal to the rank. Empirically, we show that the decomposition can reduce the size of a simple encoder-decoder Transformer to as little as 4 million parameters, including 2 million for the model and 2 million for the embeddings. The technique can be employed automatically to any matrix, including embedding, dense, attention, and output layers, without requiring any modification of the model architecture.

heading: 2 Related Work

text:
The two approaches most closely related to ours are word2ket embeddings [17] and parameterized hypercomplex multiplication (PHM) linear layers [18]. Word2ket represents input embedding matrix in language models via a sum of Kronecker products, each involving two smaller matrices. It achieves above 100-fold reduction in the number of trainable embedding parameters; however, their approach relies on lazy tensors to efficiently extract embeddings for a limited number of words in the input sequence, and does not extend to output embedding matrix or other model matrices that perform general linear transformations of their input vectors. PHM linear layers [18] arise from generalizing hypercomplex number layers [19] to arbitrary dimensionality r, with the arithmetic over the numbers learned during training. For a given dimensionality r, the arithmetic takes form of a sum of r Kronecker products. Irrespective of the original matrix size n×m, each Kronecker product involves a small r × r matrix and a larger n/r × m/r matrix, which together have r3 + nm/r trainable parameters. For r used in practice, the second term dominates, resulting in r-fold reduction in size. PHM layers with r as large as 16 have been demonstrated to be effective, leading to a reduction of a 44 million parameter Transformer to a 2.9 million parameter PHM-Transformer. These numbers do
not include the embedding matrix for token embeddings which is not reduced in size by the PHM approach; and which for example for 32, 000 BPE tokens requires about 16 million parameters.
Both of the Kronecker-based compact language model approaches [17, 18], as well as Kroneckerbased convolutional [20] and recurrent networks [21], demonstrate that a sum of Kronecker products leads to very compact representations. Here, we advance theoretical understanding of why this low-rank representation has advantages in terms of expressivity compared to a standard matrix factorization with the same rank. We also provide a more efficient, more general Kronecker-based representation. Unlike word2ket, it can be used in linear transformations anywhere in the model, an unlike PHM it also applies to embedding matrices. To represent an n×m matrix, our approach requires 2r √ mn instead of r3 + nm/r parameters required by PHM. For example, for r = 16, the highest used by PHM, for a Transformer with hidden dimension of 512, each 512×2048 feed-forward matrix in the self-attention head would be reduced from 1 million parameters to 69,632 parameters by PHM, but to 16,384 by our approach.
Kronecker product of two matrices is a concrete-basis representation of a finite-dimensional linear operator defined by a tensor product of two underlying linear operators. Thus, Kronecker-productbased representations, including PHM layers and ours, can be seen as special case of a more general family of higher-order tensor-product-based representations. Kronecker-based representation unfolds a n×m matrix from two smaller matrices, n1 ×m1 and n2 ×m2, with m = m1m2 and n = n1n2. In tensor-based representations, this is generalized to m = ∏o j=1mj and n = ∏o j=1 nj for some tensor order o, leading to unfolding the n×m matrix from an o-way tensor that is then decomposed using low rank leading to reduction in parameters. For example, the tensor-train representation [22] represents a matrix using a series of core tensors with matching ranks. Tensor decomposition has been used successfully in convolutional networks [23, 24] and recurrent networks [25]. In the context of language models, tensor-train representation has been used to construct tensorised embeddings [26], which use order-three tensor representation to obtain 60-fold reduction in embedding in a sequenceto-sequence Transformer, and also use order-six tensors to achieve almost 400-fold reduction in an LSTM model for sentiment analysis. Similarly, word2ket [17] uses order-four tensor representation to achieve more than 90,000-fold reduction in the embedding size for question-answering LSTM model. Beyond embeddings, Tensorized Transformer [27] uses block-term tensor decomposition – a combination of CP and Tucker decompositions – to reduce the multi-headed self-attention layer with nheads by a factor 1/nheads, leaving other layers intact; this allowed for reducing a 52M Transformer down to 21M trainable parameters, with little drop in BLEU scores on WMT’16 En-De translation task compared to original Transformer. So far, none of the Kronecker- or tensor-productbased approaches have been applied comprehensively to all components of a Transformer model – embeddings, attention, and feed-forward layers.

heading: 3 Limitations of Factorized Matrices in Deep Networks

text:
A simple approach for reducing the size and complexity of a model involves replacing each linear transformation n×mmatrixW with its factorized, low-rank representation involving an n×r matrix A and an r ×m matrix B, with the rank r m,n. During training and inference, the product AB is used instead of W , essentially replacing W with a two linear layers, A stacked on B. In this way, the number of parameters reduces from nm needed to represent W to r(n+m) needed to represent A and B.
One potential approach to use the factorized representation is to take an already trained model, and factorize each matrix Wj into its low-rank approximation Ŵj = AjBj using a method such as SVD. However, analyzing an idealized deep model with linear activations shows that errors from the layers add up [28], ‖ ∏ jWj − ∏ j Ŵj‖ = ∑ j‖Wj − Ŵj‖. Instead, a typical approach, used also here, is to train a model de novo in its factorized form, using randomly initialized matrices Aj and Bj , without attempting to approximate individual linear transformations of the non-factorized model.
Representing an n × m matrix W as a low-rank product AB limits the linear transformation – the output is constrained to be a low-dimensional subspace of Rn. Stacking rank-r factorized linear transformations into an idealized deep network with linear activations does not improve the expressiveness, the dimensionality of the space of possible outputs is still constrained by the smallest of the decomposed layers rank. While this may have regularizing effect for medium values of the rank, it may prevent from using very small ranks and thus from creating compact models.

heading: 4 Stacked Kronecker Product-based Representations

text:
For n = m, the standard factorized (n× r) (r ×m) → n × m representation is optimal among pairwise-product-based representations in the sense that all elements of a column of A are multiplied with each element of a row of B. For non-square matrices, for example matrix A is smaller than B, a more efficient all-pairs product representation can be made by increasing the size of A and decreasing the size of B.
In Shapeshifter, instead of the standard (n× r) (r ×m) → n ×m factorized representation, we rearrange the mn parameters of the original matrix W into a √ mn× √ mn matrix3 W ′, and then use a factorized representation of W ′ as a √ mn× r matrix A and an r ×
√ mn matrix B:(√
nm× r ) ( r × √ nm ) multiply−−−−→ √nm×√nm reshape−−−−→ √n×√m×√n×√m
transpose−−−−−→ √ n× √ n× √ m× √ m reshape−−−−→ n×m
This representation uses 2r √ mn parameters instead of nm.
We apply this representation to all matrices in attention and feed-forward layers in all multi-head attention blocks, and to the embedding matrices. For non-square matrices, the representation results in parameter saving compared to standard factorization with the same rank r. For square matrices, no parameter saving is achieved. However, as we show below, the way the matrix is decomposed allows for increased expressive power. For r = 1, the n×m matrix W arises from taking products of each element from √ nm× 1 matrix A with each element from 1 × √ nm matrix B. It is known to be equivalent [29] to a Kronecker product A⊗B involving √ n× √ m matrices A, B resulting from reshaping the single-column and single-row matrices, respectively. Beyond r = 1, the representation is equivalent to representing W as sum of r Kronecker products, W = ∑r j=1Aj ⊗ Bj . Compared to the equivalent sum of Kronecker-products representation , using the matrix factorization representation avoids the explicit summation of r terms. In current libraries, it results in a simpler implementation involving only matrix multiplication, reshaping and transposing, and is more efficient compared to performing individual Kronecker products followed by an explicit summation along the rank dimension.

heading: 4.1 Expressiveness of stacked Kronecker-product Layers

text:
Certain types of matrices can easily be expressed as a sum of Kronecker products, but not in a r-rank factorized matrix form. For example, identity mapping on Rn does not admit low-rank factorization, but it can be decomposed as a single Kronecker product of two smaller identity matrices. More broadly, the output of a sum of Kronecker products layer is not limited to r-dimensional subspace; a Kronecker product of two orthogonal matrices is an orthogonal matrix [30], thus even a single A⊗B term can model isomorphisms Rn → Rn. The converse is not true, most matrices require Kronecker representation with a high rank, and the question of approximating various types of matrices with a sum of Kronecker products has received ample attention [31, 32, 30, 33].
Here, we focus on a different question – how does the expressive power of a stack of layers, each involving a sum of Kronecker products instead of generic linear transformation, grow with the network depth. We analyze an idealized network formed by a stack of Kronecker layers with linear activations. Given high-enough rank, a single layer is enough to represent any matrix. We first analyze what the depth needs to be if rank is restricted to two, and then investigate how quickly the depth can be reduced as the rank increases.
Below, we use the following notation. Support supp(x) of vector x is a set of indices at which x is not null. We use [n] = {1, ..., n}. I is identity matrix. 1jk is a matrix with unity at row j and column k and null elsewhere; 1k is a shorthand for 1kk. To avoid ambiguity, we use square brackets for individual entries of a matrix or a vector; thus A[j, :] represents j-th row of matrix A, while Aj represents j-th matrix out of a series, as used above. For index set Z of cardinality k and a square matrixA, byA[Z] we represent a k×k submatrix ofA obtained by taking the intersection of columns and rows from set Z.
We first define k-variant matrices, which will be a tool in exploring stacked layers. 3If √ mn is not an integer, we round it up, and use a submatrix as representation of W .
Definition 1 (k-variant matrix). For an n × n orthogonal matrix U and an index set Z ⊂ [n], by writing UZ we indicate that both U and its inverse UT act as identity on all the dimensions not in set Z. That is, for each x ∈ Rn, supp(Ux− x) ⊆ Z and supp(UTx− x) ⊆ Z. If Z has cardinality at most k, we call UZ a k-variant matrix.
The notion of k-variant matrices generalizes Givens rotation matrices [34] used for example in QR decomposition. A Givens rotation in n dimensions is any rotation that acts on a plane spanned by two of the n coordinate axes – it is thus 2-variant. QR decomposition can be used to show that each matrix can be represented as a product of 2-variant matrices. We will then show that 2-variant matrices, or more broadly, k-variant matrices for k ≤ √ n, can be represented using sum of Kronecker products.
Theorem 2 (Decomposition into layers involving Kronecker products). Let n = m2 for somem ∈ N. Any orthogonal n× n matrix U can be represented as
U = L∏ i=1 2∑ j=1 Aij ⊗Bij ,
where Aij and Bij are √ n× √ n matrices, and L = O ( n2 ) .
Proof. QR decomposition using Givens matrices will result in Q being a product of up to n(n− 1)/2 Givens rotations, which are 2-variant, and R an upper triangular matrix which for orthogonal matrices is a diagonal matrix with +1 and −1 elements. The sign of any pair of −1 diagonal entries can be negated by a 2-variant matrix; at most n/2 such matrices are needed to convert the diagonal into identity, since there are at most n− 1 negative elements in R. In total, up to n2/2 2-variant matrices are needed to represent any given orthogonal matrix.
Next, we set to show that any 2-variant matrix can be represented as a product of at most three matrices, each of the form A1 ⊗B1 +A2 ⊗B2, using m×m matrices.
Let U{k,q} be an arbitrary 2-variant orthogonal m2 ×m2 matrix for arbitrary k, q ∈ [m2]. That is, y = U{k,q}x acts only on x[k] and x[q] to produce y[k] and y[q], for l 6= k, q we have y[l] = x[l]. Consider two vectors, xA, xB ∈ Rm. Kronecker product of x = xA ⊗ xB of these two vectors is a vector x ∈ Rn with entries defined through products
x[k] = xA[κ]xB [λ] for k = (κ− 1)m+ λ, x[q] = xA[π]xB [ρ] for q = (π − 1)m+ ρ
for κ, λ, π, ρ ∈ [m].
A matrixA⊗B acts on vectors x = xA⊗xB as (A⊗B)x = (AxA)⊗(BxB). For example, letB{λ,ρ} be a 2-variant matrix action on λ, ρ, then 1ν ⊗B{λ,ρ} is a acting on dimensions k = (ν − 1)m+ λ and q = (ν − 1)m+ ρ in the same way as B{λ,ρ} acts on dimensions λ and ρ, and producing either identity or null elsewhere, as shown below:
[ 0
1 0
] ⊗ [ α −β
1 β α
] =  0 . . . α −β 1 β α . . .
0
 .
First, consider k, q such that κ = π = ν for some ν ∈ [m]. Define a 2-variant B{λ,ρ} such that B[{λ, ρ}] = U [{k, q}], that is, B acts on λ, ρ in the same way as U acts on k, q. We then have
U{k,q} = 1ν ⊗B{λ,ρ} + (I − 1ν)⊗ I.
A symmetric construction holds if λ = ρ = ν. In both cases, we can represent a 2-variant orthogonal matrix by A1 ⊗B1 +A2 ⊗B2.
The general case of U{k,q} for κ 6= π and λ 6= ρ can be addressed by converting it into the case κ = π via transforming κ into π, applying the special case, and then transforming π back into κ.
We show that an orthogonal matrix V of the form V = C ⊗D + C ′ ⊗D′ exists such that U{k,q} = V T ( 1π ⊗B{λ,ρ} + (I − 1ν)⊗ I ) V,
where B[{λ, ρ}] = U [{k, q}], the λ, ρ columns/rows of B are formed by taking submatrix of U defined by columns/rows k, q.
To construct V , set C = Pκ→π, a permutation matrix that moves κ to π, and D = 1λ. Also set C ′ = I and D′ = I − 1λ. Columns of C ⊗D and C ′ ⊗D are linearly independent, and each of the two matrices is orthogonal, thus V is orthogonal.
Consider x = xA ⊗ xB , we then have x′ = V x = Pκ→πxA ⊗ 1λxB + IxA ⊗ (I − 1λ)xB .
The first term will use the permutation to move xA[κ] to position π, where it will be multiplied by xB [λ], and placed at x′[k′] for k′ = (π − 1)m+ λ; all other entries xA[·]xB [λ] will be rearranged as well by the permutation. For m = 3, λ = 2, κ = 1, π = 2, the transformation that moves x[k = 2] = x[(κ − 1)m + λ] = α1β2 into x′[k′ = 5] = x′[(π − 1)m + λ] and thus brings it into position ready for the special case, is illustrated below[
0 1 0 1 0 0 0 0 1 ][ α1 α2 α3 ] ⊗ [ 0 0 0 0 1 0 0 0 0 ][ β1 β2 β3 ] = [ α2 α1 α3 ] ⊗ [ 0 β2 0 ] =
= [0 α2β2 0 0 α1β2 0 0 α3β2 0] T .
The second term will preserve xA[·]xB [·] for all entries of xB except λ; this includes ρ, since in the general case λ 6= ρ. The k, q entries of x will be now at k′, q entries of x′, and applying 1π ⊗ B{λ,ρ} + (I − 1ν)⊗ I on x′ will be equivalent to acting on positions k, q of x. The inverse orthogonal transformation V T will restore the transformed values at positions k′, q, as well as all other entries xA[·]xB [λ], back to their original positions. The transformation V T is of the same form as V , except using P−1κ→π instead of Pκ→π , since the other matrices involved in V are diagonal with unit-magnitude entries. In summary, in the general case, m×m matrices Aij and Bij exist such that any 2-variant orthogonal matrix can be represented as ∏3 i=1 ∑2 j=1Aij ⊗Bij .
Together, any orthogonal n = m2 dimensional matrix U can be decomposed into U = ∏3n2/2 i=1 ∑2 j=1Aij ⊗Bij .
Corollary 3. Any linear n ×m matrix W can be represented as a product of a finite sequence of sums of two Kronecker products of dne × dme matrices followed by selecting an n×m submatrix.
Proof. The result immediately follows by considering an orthogonal matrix with dimension max(dne2 × dme2), observing that an arbitrary, possibly rank-deficient diagonal matrix D can be decomposed as a product of a series of matrices of the form 1k⊗(I−(D[j, k]−1)1j)+(I−1k)⊗I , one per diagonal entry D[j, k], which allows for producing arbitrary rank-deficient square matrices, from which non-square matrices can be selected.
The above theorem shows that stacking layers represented compactly as ( √ nm× r) (r × √ nm) for r = 2 increases the expressive power of the network. In the idealized case with linear activation function, adding such layers allows the network to eventually produce arbitrary Rm → Rn linear transformations, unlike stacking standard factorized layers (n× r) (r ×m), which always maps into an r-dimensional subspace irrespective of the number of layers.
With rank r = 2, we may need as many as 3n2/2 layers to represent an orthogonal n× n matrix. On the other extreme, only one layer is needed if the rank is n, since arbitrary n× n matrix can be seen as a concatenation of n tiles of size √ n× √ n. As long as we define n indicator √ n× √ n matrices 1jk and use them as tile position selectors, we can use a sum of Kronecker products of the selection with the corresponding tile to produce the matrix. In between these extremes, the number of layers needed to increase the expressiveness of the network to allow arbitrary transformations depends on the rank r in the Shapeshifter representation in the following way.
Theorem 4 (Depth-rank tradeoff for deep models involving Kronecker products). Let n = m2 for some m ∈ N, and let r ≤ m. Any orthogonal n× n matrix U can be represented as a product of a sequence of sums of two Kronecker products of m×m matrices,
U = L∏ i=1 r∑ j=1 Aij ⊗Bij ,
where Aij and Bij are √ n× √ n matrices, and L = O ( n2/r ) .
Proof. Instead of 2-variant orthogonal matrices, we use r-variant ones. Consider r = 3, with U{k,q,t}, and with k = (κ− 1)m+ λ, q = (π − 1)m+ ρ, and t = (τ − 1)m+ υ. We now have two groups of indices, κ, π, τ and λ, ρ, υ. If κ = π = τ = ν for some ν, we have a special case as in Theorem 2
U{k,q,s} = 1ν ⊗B{λ,ρ,υ} + (I − 1ν)⊗ I;
we have a similar special case if λ = ρ = υ.
The general case involves no equality in either of the two groups. Then, we can proceed similarly as for 2-variant general case, but we need a permutation matrix Pκ→τ and separately a permutation matrix Pπ→τ , which allow us to define
x′ = V x = Pκ→τxA ⊗ 1λxB + Pπ→τxA ⊗ 1ρxB + IxA ⊗ (I − 1λ − 1ρ)xB .
Here, V is an orthogonal transformation that uses a sum of three Kronecker products as a transformation that maps the general case to the κ = π = τ case. The construction naturally extends into higher r as long as r ≤ m, the number of diagonal entries in matrices A: we need V to be a sum of r Kronecker products involving distinct dimensions of A. In summary, we need less than∏3 i=1 ∑r j=1Aij ⊗Bij to represent arbitrary r-variant orthogonal matrix. A product of two k-variant matrices is at most 2k-variant. Grouping r/2 Givens rotations or, more generally, 2-variant matrices can result in at most a r-variant matrix, thus we need at most L = O ( 3n2/r ) layers, each involving a Kronecker product of rank at most r.
As is common in theoretical work on deep networks, the results are limited to idealized multilayer models with linear activations. However, they can provide insights into networks with nonlinearities; for example, activations such as ReLU that are linear on large part of their input will lead to piece-wise linear network, and the results above hold in a piece-wise fashion.

heading: 5 Experimental Results

text:
We validated the approach on machine translation using sequence-to-sequence models. The experiments were performed on a single V100 (longest run: 6 days) or A100 GPU (longest run: 1 day). Time overhead experiments were performed on a dedicated workstation with a single NVIDIA RTX 3090 GPU. As the basis for Shapeshifter, we used Transformer [1], with embedding and encoder/decoder layers replaced with compact representations. In Shapeshifter, we used small rank for all matrices in the multi-head self-attention blocks. For embedding matrices, which are much larger and can be reduced more effectively, we used higher rank, while aiming to keep the total size of factorized embeddings below the total size of the rest of the factorized encoder/decoder. The time overhead
(see Table 1) introduced by the representation is modest, up to 17% increase in running time during training, and up to 13% during predictions with a trained model.
We compared our approach to two recently proposed approaches for reducing model size: DeLighT [14] and PHM Layers [18]. DeLighT uses separate vocabulary for the source language encoder embeddings, and a different vocabulary for the target language decoder embedding and output layer. PHM Layers, like its predecessor Quaternion Transformer [19], uses the same Byte-Pair Encoding [38] in both encoder and decoder. To facilitate comparisons, we followed both approaches and trained two Shapeshifter variants, one with separate embeddings and one with single embedding.

heading: 5.1 Comparison with PHM Layers

text:
In comparisons with PHM layers [18], we used T5-small Transformer [35], which shares a single embedding matrix in both the encoder and the decoder. We use Huggingface Transformers [40] for PyTorch [41] implementation. We evaluate the approach on IWLSLT’14 German-to-English (De-En) [37] and WMT’18 English-to-Romanian (En-Ro) [36] datasets. For the first dataset we use learning
rate 2e-3 with batch size of 128, while for the larger En-Ro dataset we use 3e-3 with batch size of 192. We used LAMB optimizer [42] with inverse square root scheduler, dropout 0.1, no weight decay, and 0.1 label-smoothed cross entropy loss.
To facilitate comparison of our results with PHM Layers approach under the same conditions, we reimplemented the approach and applied it to T5-Small Transformer that we used as basis for experiments with our method. The small increase in size of the PHM Transformer observed in our re-implementation of PHM layers compared to the parameter size provided in the literature [18] (3.2M vs 2.9M parameters) results from the fact that the underlying T5 Transformer that we build on uses three separate dhid × dhid matrices for self-attention query, key, and value weights for all heads, while the Transformer used in [18] stores all three as a single, concatenated dhid × 3dhid matrix. The results presented in Table 2 and in the left panel of Figure ?? indicate that while our full models are four times smaller than full PHM models, they offer very similar on-task performance. The parameter reduction comes not only from factorized embeddings in Shapeshifter, but also from matrices inside the attention blocks in encoder/decoder: both Shapeshifter 16 / 256 and PHM n = 16 use a sum of 16 Kronecker products to represent a matrix, yet PHM results in a 1.5 times larger encoder/decoder when applied to the same exact base architecture.

heading: 5.2 Comparison with DeLighT

text:
Following DeLighT [14], we used Transformer models as implemented in fairseq [45]. We benchmarked the Shapeshifter models on four datasets: IWSLT’14 German-to-English (De-En) [37], WMT’16 English-to-Romanian (En-Ro) [43], WMT’14 English-to-German (En-De) [44], and WMT’14 English-to-French (En-Fr) [44]. The specific Transformer architecture used by DeLighT, and also in our experiments, varies depending on the dataset. For the WMT’14 En-De we used 12 transformers blocks with 4 attention heads, 512 as the embedding dimension and 1024 for the fan out fully-connected layer. For the other three datasets we used 12 transformers blocks with 8 attention heads, 512 as the embedding dimension and 2048 for the fan out fully-connected layer. For the training setup we have used Adam optimizer, learning rate of 5e-4, inverse square root scheduler with 15K warmup steps, dropout 0.3, weight decay of 1e-4, and 0.1 label-smoothed cross entropy loss.
Compared to the smallest of the DeLighT models available for each dataset, as summarized in Table 3 and in the right panel of Figure ??, the Shapeshifter models have between 1.5 and 4 times
fewer parameters, yet offer similar quality; Shapeshifter achieves higher score than DeLighT on two datasets, and lower score on two other datasets.

heading: 5.3 Comparison with Standard Low-rank Factorization

text: To validate experimentally the theoretical results from Section 4.1 concerning the increased expressiveness of the factorization type used in Shapeshifter compared to standard n× r, r ×m low-rank factorization, we trained models factorized this way using the same encoder/decoder rank of 16, and the same embedding rank of 256 as in Shapeshifter. To test only the stack of Transformer layers, we also trained models that do not factorize the embedding matrix. We used one small dataset, IWSLT’14 German-to-English, and one large, WMT’18 English-to-Romanian. The results in Table 4 show that the size of the model is reduced substantially in the Kronecker-based approach compared to standard low-rank representation. The reduction comes predominantly from the embedding, which involve non-square matrices. Despite more than two-times larger number of parameters, standard low-rank representation shows much higher loss on the training set, with no significant difference in terms of generalization outside of the training set, confirming that the difference between the two representations results from increased expressiveness of a stack of low-rank Kronecker-product-based layers. Similar behavior can be observed when the comparison involves models where the embedding matrices are not factorized.

heading: 6 Conclusion

text:
Our theoretical and experimental results show that deep models composed of stacked low-rank Kronecker-product-based representations are more expressive, yet smaller, than equivalent models that use standard low-rank factorized matrix representations. These observations help solve a puzzle of how recent methods such as word2ket [17] and PHM layers [18] achieve their impressive parameter reduction rates for embeddings and for encoder/decoder layers, respectively, despite training models de novo, without the help of an existing large model serving as a blueprint for compression or pruning, or as a teacher for knowledge distillation. Building on these results, we provide a comprehensive Shapeshifter4 approach that works both for the embeddings and for the encoder/decoder, using more effective way to exploit Kronecker products for factorized representations. Experimental comparisons with state-of-the-art model reduction approaches, not only PHM but also recently proposed DeLighT [14], on a range of machine translation problems, show that the approach offers similar translation quality with much fewer parameters.
Funding Disclosure and Competing Interests
This work was partially funded by NSF grant IIS-1453658 to T.A. The authors have no competing interests.
4Code available at: https://github.com/tarodz/shapeshifter.

title: Shapeshifter: a Parameter-efficient Transformer using Factorized Reshaped Matrices


summaries: The paper presents a way to incorporate the tree-structure information into a transformer architecture to better learn representation for code. The tree structures are first dispatched as leaf-to-leaf "relative paths" and leaf-to-root "absolute paths", and then embedded and modulated as additional learnable parts of the query, key, and value components within the transformer attention. The authors evaluate their proposal on the task of extreme code summarization -- predicting function names. Results show improvements over 4 previous non-transformer and transformer models on 3 out of 4 programming languages tested. The authors also carry out ablation studies to understand the effect of absolute paths and pointer networks but many important questions remain unanswered.

summaries: This paper introduces a new neural network architecture for code, TPTrans (Tree Path Transformer), by modifying the attention calculation in the Transformer architecture to use information about paths in the syntax trees of programs. The modification works by computing path embeddings between pairs of leaf nodes ("relative path embeddings") and between leaf nodes and the root (forming nodes' "absolute path embeddings"), and incorporating these into the Transformer's attention module. The paper evaluates TPTrans, including baselines and variants, on code summarization -- the task of predicting a function's name given its body -- for four different programming languages. The main experiments reveal the value of including embeddings of sequences (the paths) in the attention calculation over baselines, and subsequent experimentation compares the impact of relative and absolute paths, and the effect of pointer network outputs on the model.

summaries: This work continues on earlier advances in how positional and relational encoding is used by the self attention primitive of Transformers, in particular in the context of source-code understanding. Prior work has biased self attention to encode relative relational information between inputs (e.g., edges or distance on some graph), and to separate the self-attention term due to position from that due to input.  For the former, TPTrans runs the path between any two tokens (specifically, the path in the code Parse Tree between two leaves -- the authors call this the "relative path") through an RNN, and uses the resulting embedding to bias the self attention between those tokens.  For the latter, TPTrans-$\alpha$ runs the path from the parse-tree root to a token through an RNN (the "absolute path"), and uses that to construct a separate positional self-attention term.  The results show that both TPTrans and TPTrans-$\alpha$ outperform prior work (most notably, Code Transformer), but TPTrans seems to dominate TPTrans-$\alpha$ with a smaller set of trainable parameters.

summaries: The paper proposes an Transformer-based architecture for embedding source code snippets that integrates encoding of absolute and relative AST paths for all code tokens into self-attention. Following up on prior work in relative position encoding, they map AST paths into the position encoding framework, modify the self-attention computation accordingly, and evaluate on method summarization. The approach is compared to recent Transformer-based baselines on the standard method summarization benchmark, as well as to code2seq, which pioneered path-only source code encoding. The results improve performance significantly on almost all benchmarks.

summaries: The authors present a technique for modeling source code. They point out that modeling code is different from language because of the tight binding between ASTs, data, and control flow. Also they mention some weakness of GNNs. To compete with GNNs, they introduce a new method which is a sequence-based transformer that incorporates graph structure through the attention weights. The model is shown to be effective for code summarization.

abstractText: Learning distributed representation of source code requires modelling its syntax and semantics. Recent state-of-the-art models leverage highly structured source code representations, such as the syntax trees and paths therein. In this paper, we investigate two representative path encoding methods shown in previous research work and integrate them into the attention module of Transformer. We draw inspiration from the ideas of positional encoding and modify them to incorporate these path encoding. Specifically, we encode both the pairwise path between tokens of source code and the path from the leaf node to the tree root for each token in the syntax tree. We explore the interaction between these two kinds of paths by integrating them into the unified Transformer framework. The detailed empirical study for path encoding methods also leads to our novel state-of-theart representation model TPTrans, which finally outperforms strong baselines. Extensive experiments and ablation studies on code summarization across four different languages demonstrate the effectiveness of our approaches. We release our code at https://github.com/AwdHanPeng/TPTrans.

authors:
Abram Hindle
Earl T Barr
Mark Gabel
Zhendong Su
Premkumar Devanbu

title: On the naturalness of software

venue: Communications of the ACM,

authors:
Vincent J Hellendoorn
Premkumar Devanbu

title: Are deep neural networks the best choice for modeling source code

venue: In Proceedings of the 2017 11th Joint Meeting on Foundations of Software Engineering,

authors:
Bolin Wei
Ge Li
Xin Xia
Zhiyi Fu
Zhi Jin

title: Code generation as a dual task of code summarization

authors:
Uri Alon
Shaked Brody
Omer Levy
Eran Yahav

title: code2seq: Generating sequences from structured representations of code

venue: arXiv preprint arXiv:1808.01400,

authors:
Sepp Hochreiter
Jürgen Schmidhuber

title: Long short-term memory

venue: Neural computation,

authors:
Seohyun Kim
Jinman Zhao
Yuchi Tian
Satish Chandra

title: Code prediction by feeding trees to transformers

venue: arXiv preprint arXiv:2003.13848,

authors:
Miltiadis Allamanis
Marc Brockschmidt
Mahmoud Khademi

title: Learning to represent programs with graphs

venue: arXiv preprint arXiv:1711.00740,

authors:
Yaqin Zhou
Shangqing Liu
Jingkai Siow
Xiaoning Du
Yang Liu

title: Devign: Effective vulnerability identification by learning comprehensive program semantics via graph neural networks

authors:
Patrick Fernandes
Miltiadis Allamanis
Marc Brockschmidt

title: Structured neural summarization

venue: arXiv preprint arXiv:1811.01824,

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

title: Attention is all you need

venue: arXiv preprint arXiv:1706.03762,

authors:
Vincent J Hellendoorn
Charles Sutton
Rishabh Singh
Petros Maniatis
David Bieber

title: Global relational models of source code

venue: In International conference on learning representations,

authors:
Peter Shaw
Jakob Uszkoreit
Ashish Vaswani

title: Self-attention with relative position representations

venue: arXiv preprint arXiv:1803.02155,

authors:
Daniel Zügner
Tobias Kirschstein
Michele Catasta
Jure Leskovec
Stephan Günnemann

title: Language-agnostic representation learning of source code from structure and context

venue: arXiv preprint arXiv:2103.11318,

authors:
Zhilin Yang
Zihang Dai
Yiming Yang
Jaime Carbonell
Ruslan Salakhutdinov
Quoc V Le

title: Xlnet: Generalized autoregressive pretraining for language understanding

authors:
Uri Alon
Roy Sadaka
Omer Levy
Eran Yahav

title: Structural language models of code

venue: In International Conference on Machine Learning,

authors:
Hoa Khanh Dam
Truyen Tran
Trang Pham

title: A deep language model for software code

venue: arXiv preprint arXiv:1608.02715,

authors:
Song Wang
Devin Chollak
Dana Movshovitz-Attias
Lin Tan

title: Bugram: bug detection with n-gram language models

venue: In Proceedings of the 31st IEEE/ACM International Conference on Automated Software Engineering,

authors:
Miltiadis Allamanis
Hao Peng
Charles Sutton

title: A convolutional attention network for extreme summarization of source code

venue: In International conference on machine learning,

authors:
Srinivasan Iyer
Ioannis Konstas
Alvin Cheung
Luke Zettlemoyer

title: Summarizing source code using a neural attention model

authors:
Lili Mou
Ge Li
Lu Zhang
Tao Wang
Zhi Jin

title: Convolutional neural networks over tree structures for programming language processing

venue: In Thirtieth AAAI Conference on Artificial Intelligence,

authors:
Uri Alon
Meital Zilberstein
Omer Levy
Eran Yahav

title: code2vec: Learning distributed representations of code

venue: Proceedings of the ACM on Programming Languages,

authors:
Richard Socher
Cliff Chiung-Yu Lin
Andrew Y Ng
Christopher D Manning

title: Parsing natural scenes and natural language with recursive neural networks

venue: In ICML,

authors:
Qian Chen
Xiaodan Zhu
Zhenhua Ling
Si Wei
Hui Jiang
Diana Inkpen

title: Enhanced lstm for natural language inference

venue: arXiv preprint arXiv:1609.06038,

authors:
Yikang Shen
Shawn Tan
Alessandro Sordoni
Aaron Courville

title: Ordered neurons: Integrating tree structures into recurrent neural networks

venue: arXiv preprint arXiv:1810.09536,

authors:
Vighnesh Shiv
Chris Quirk

title: Novel positional encodings to enable tree-based transformers

venue: Advances in Neural Information Processing Systems,

authors:
Yau-Shian Wang
Hung-Yi Lee
Yun-Nung Chen

title: Tree transformer: Integrating tree structures into self-attention

venue: arXiv preprint arXiv:1909.06639,

authors:
Zihang Dai
Zhilin Yang
Yiming Yang
Jaime Carbonell
Quoc V Le
Ruslan Salakhutdinov

title: Transformer-xl: Attentive language models beyond a fixed-length context

authors:
Colin Raffel
Noam Shazeer
Adam Roberts
Katherine Lee
Sharan Narang
Michael Matena
Yanqi Zhou
Wei Li
Peter J Liu

venue: arXiv preprint arXiv:1910.10683,

authors:
Guolin Ke
Di He
Tie-Yan Liu

title: Rethinking the positional encoding in language pre-training

venue: arXiv preprint arXiv:2006.15595,

authors:
Antoine Bordes
Nicolas Usunier
Alberto Garcia-Duran
Jason Weston
Oksana Yakhnenko

title: Translating embeddings for modeling multi-relational data

venue: Advances in neural information processing systems,

authors:
Kyunghyun Cho
Bart Van Merriënboer
Dzmitry Bahdanau
Yoshua Bengio

title: On the properties of neural machine translation: Encoder-decoder approaches

venue: arXiv preprint arXiv:1409.1259,

authors:
Hamel Husain
Ho-Hsiang Wu
Tiferet Gazit
Miltiadis Allamanis
Marc Brockschmidt

title: Codesearchnet challenge: Evaluating the state of semantic code search

authors:
Diederik P Kingma
Jimmy Ba

title: Adam: A method for stochastic optimization

venue: arXiv preprint arXiv:1412.6980,

heading: 1 Introduction

text:
Machine learning for source code aims at building models that learn semantic embedding of programs. The initial representation of source code relied on sequential models adopted from natural language processing, such as n-gram language model [Hindle et al., 2016, Hellendoorn and Devanbu, 2017], Recurrent Neural Networks (RNNs) [Wei et al., 2019] and so on. However, source code is more logical than natural languages, rich in structured information such as Abstract Syntax Tree (AST). Therefore, these previous works struggle to capture the structural complexity of source code.
Some works leverage the program AST to model source code structure and linearize the graph by traversing it. Code2seq [Alon et al., 2018] represents source code as a set of pairwise paths over AST where each path is compressed to a vector using LSTMs [Hochreiter and Schmidhuber, 1997]. Code2seq obtains state-of-the-art for code representation using only AST information, demonstrating the effectiveness of path encoding. [Kim et al., 2020] leveraged another kind of path from the leaf node to the root of AST by traversing up its ancestors and then coupled the representation of paths with the source code.
On the other hand, several works also leverage structured graph neural networks for modelling source code directly. [Allamanis et al., 2017] used Gated Graph Neural Network (GGNN) to embed the semantically meaningful relationships in the code. [Zhou et al., 2019] proposed a novel graph neural network with composite programming representation for vulnerability identification. However, GNNs
∗Corresponding author
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
typically rely on synchronous message passing, requiring message passing iterations to aggregate information [Allamanis et al., 2017, Fernandes et al., 2018]. Besides, primarily for computational reasons, GNNs for programs usually compute few message-passing iterations. For these reasons, the representation of GNNs tends to be local and struggle to leverage long-range interactions.
On the contrary, Transformer [Vaswani et al., 2017] allows global information flow at each step but lacks the ability to model the structural complexity of source code. Recently, some works have explored introducing structural inductive bias in Transformer to access the global structural representation of source codes. [Hellendoorn et al., 2019] proposed Graph Relation Embedding Attention Transformer (GREAT), which extends [Shaw et al., 2018] and biases Transformer with relational information from graph edge types. After that, [Zügner et al., 2021] proposed Code Transformer based on XLNet [Yang et al., 2019], which computes pairwise distances on AST and integrate multiple relations into the attention module.
In this paper, we pursue the research line of combining Transformer with additional structure information of source code. Our first starting point is the structural model of Code2seq [Alon et al., 2018], which obtains state-of-the-art of code summarization using only pairwise path information in AST. However, Code2seq lacks the modelling of context, which leads us to explore combining path representation with source code context. The other starting point is Code Transformer [Zügner et al., 2021], which counts node distances cross AST to capture source code structure. However, the different nodes combinations of paths contain plenty of structure information, which is overlooked by only encoding distances.
To overcome the drawback shown in previous works, we adopt the ideas of encoding path to represent source code and integrate them to bias the attention module of Transformer. Specifically, we encode the path between source code tokens across AST and the path from leaf node to tree root for each token. Then we draw inspiration from both relative and absolute positional encoding methods in the NLP field and modify them to integrate path encodings into Transformer. These paths introduce inductive bias into the attention module of Transformer, powering it to know the structure of source codes. For clarity, we name the first path as the relative path and the last one as the absolute path (see Fig1), since much similarities between path and position on relative or absolute encoding exist.
Intuitively, the relative path represents the structural relationship between tokens and shows specific patterns of code block such as For loop, etc. Meanwhile, the absolute path complements the structural information for each token and reveals the program behaviour on it, such as Assignment and Call, etc. Both two kinds of paths contain plenty of structural information of code syntax. Actually, lots of works have introduced these paths into many novel models [Kim et al., 2020, Alon et al., 2018, 2020]. However, to our best knowledge, the feature relationship between these two paths for code representation learning is still not completely studied. In this paper, we integrate these paths into the unified Transformer framework to analyze each effectiveness and their relationship. The detailed
empirical study also leads to our novel code representation model TPTrans, which means encoding the Tree Path into the Transformer. We show the effectiveness of our approaches on the code summarization task across four different languages datasets, in which the model predicts a method’s name based on its body. Our model finally outperforms solid baselines and obtains state-of-the-art on most datasets. The contributions of this paper are summarized as follows:
1. We propose TPTRANS which integrates path encoding in Transformer, powers it to know the structure of source code, and our model significantly outperforms existing baselines.
2. We present the empirical study to investigate the relationship between two kinds of path encoding methods proposed in previous works, shedding light on this line of future work.

heading: 2 Related work

text:
Representation learning for source code The growing availability of open-source code creates opportunities for machine learning from big code. Early research works learned language model on raw text data for code [Dam et al., 2016, Wang et al., 2016, Allamanis et al., 2016, Iyer et al., 2016], providing evidence for the naturalness assumption [Hindle et al., 2016]. These models treat codes as sequences rather than trees and struggle to model the structural complexity of source code. After that, some works leveraged structural information to model source code. [Mou et al., 2016] proposed tree-based convolutional neural networks to represent source code. [Alon et al., 2018, 2019] treated source code as weighted combination of pairwise paths in abstract syntax trees. Several works also use structured graphical models to represent structural aspects of source code directly. For example, [Allamanis et al., 2017] proposed GGNN to represent program graphs consisting of AST together with control-flow and data-flow edges. After that, [Hellendoorn et al., 2019, Zügner et al., 2021] proposed to bias the self-attention computation of Transformer given the underlying graph structure. In this paper, we explore integrating path representation of syntax tree into Transformer, powering it to capture the structure of code snippets.
Structural representation of languages It has been proved that modelling the language structure helps to represent the compositional aspects of natural language [Socher et al., 2011] and improves the generalization of models [Chen et al., 2016]. Several works [Shen et al., 2018, Shiv and Quirk, 2019] proposed new architectures for modeling tree structures, including in the domain of natural languages processing [Wang et al., 2019] and code representation learning [Hellendoorn et al., 2019, Zügner et al., 2021]. We pursue this research line and show the benefits of integrating tree structure as inductive bias.
Positional encoding for Transformer The absolute positional encoding used in Transformer has been proved inefficient to model the order information of sequence. [Shaw et al., 2018] firstly proposed to leverage relative positional encoding to Transformer. After that, Transformer-XL [Dai et al., 2019] re-parameterized the self attention of Transformer to integrate relative positional encoding and T5 [Raffel et al., 2019] simplified the vector representation of relative positions in [Shaw et al., 2018] to scalars. Recently, several works also tried to enhance positional encodings in Transformer to integrate structural inductive bias. For example, [Shiv and Quirk, 2019] extended the sequential positional encoding to tree-based positional encoding in Transformer, and [Zügner et al., 2021] modified the relative positional encoding of Transformer-XL to incorporate multiple relations in tree structural. In this paper, we adopt the idea to modify positional encodings of Transformer and introduce structural bias of path information in the syntax tree for self-attention computing.

heading: 3 Approaches

text:
Self-attention is one of the key components of Transformer. The attention module can be formulated as querying a dictionary with key-value pairs. The matrix form equation for attention with a single head is
Attention(Q,K, V ) = softmax( QKT√
d )V, (1)
where d is the dimension of the hidden representations, and Q(Query), K(Key), V (Value) are specified as the hidden representations of the previous layer. For simplicity, we omit the index of
layer, denote x = (x1, x2 · · · , xn) and z = (z1, z2 · · · , zn) as the input and output of self-attention module in the same layer respectively, where n is the sequence length. Then Eq.1 can be rewritten as
zi = n∑ j=1 exp (αij)∑n j′=1 exp (αij′) (xjW V ), where αij = 1√ d (xiW Q)(xjW K)T , (2)
where WQ ∈ Rd×dk , WK ∈ Rd×dk , WV ∈ Rd×dv and we set dk = dv = d. Eq.2 is oblivious to structured input since it views input as unordered vector sets. In the NLP field, several works [Shaw et al., 2018, Ke et al., 2020, Raffel et al., 2019, Dai et al., 2019] propose different ways to bias Transformer towards sequential inputs by incorporating positional information into the self-attention module. Generally, there are two categories of positional encoding methods for self-attention: absolute positional encoding and relative positional encoding.
The original Transformer [Vaswani et al., 2017] use absolute positional encoding to represent positions. Given a sentence, a real-valued vector pi ∈ Rd will be added to the word embedding wi at position i, and the sum vector will be used as the input to Transformer. [Ke et al., 2020] further proposed TUPE to untie the correlations between absolute positions and words and directly model the relationship between a pair of words or positions using different projection matrices, that is
αij = 1√ 2d (xiW Q)(xjW K)T + 1√ 2d (piU Q)(pjU K)T , (3)
where UQ, UK ∈ Rd×d are projection matrices for positional embedding and scaling term √ 2d is used to retain the magnitude of αij .
Using different pi helps Transformer distinguish the words at different positions i. However, the absolute positional encoding is not effective in capturing the relative word orders. Therefore, [Shaw et al., 2018] proposed the relative positional encoding, that is
zi = n∑ j=1 exp (αij)∑n j′=1 exp (αij′) (xjW V + aVij), where αij = 1√ d (xiW Q)(xjW K + aKij ) T (4)
where aVij , a K ij ∈ Rd are learnable parameters and can be viewed as the relative position embedding between position i and j. Eq.4 extends self-attention to consider the pairwise positional relationship and models it as a directed edge between head and tail word entities [Shaw et al., 2018], which means the tail entity should be close to the embedding of the head entity plus vector depending on the relationship [Bordes et al., 2013].
We draw inspiration from the ideas of encoding absolute and relative positional information and modify Eq.3 and Eq.4 to incorporate path information in the syntax tree of source code.

heading: 3.1 Integrating path encodings in attention module

text:
A syntax tree uniquely represents a source code snippet in the given language and grammar. In a syntax tree, the leaves are called terminals, and the non-leaf nodes are called nonterminals. [Alon et al., 2018, 2019] used AST to represent a source code snippet. They considered pairwise paths between terminals, representing them as sequences of the terminal and nonterminal nodes in AST. [Kim et al., 2020] used another kind of path from the terminal to AST root by traversing up its ancestors, and then coupled paths representation with tokens and fed the sum representation into Transformer.
We adopt the above methods of leveraging paths and integrate them into the self-attention module of the unified Transformer framework. We name the path between terminals as the relative path and the path from terminal to root as the absolute path. Technically, we modify the positional encoding methods in the NLP field to integrate path encoding into Transformer, which shows natural analogies from the positional to path encoding ways as follows.
All kinds of paths are composed of nonterminals from a limited vocabulary of syntax grammar. We use a learned embedding matrix to convert nonterminals to vectors first.
Relative path encoding The embed vectors of a relative path are first fed into a sequence encoder. We use a bi-directional GRU [Cho et al., 2014] to encode the path and use the final state to represent it, that is rij = GRUr(Pathxi→xj ), (5) where Pathxi→xj is the nonterminal vector sequence of the relative path between tokens xi and xj , and rij ∈ Rd is the final vector representation of this path. Specifically, Pathxi→xj is a vector list [n0, n1, . . . , nm], where m is the path length and each n is the embedding vector of each nonterminal node which is looked up from the embedding matrix of node type.
The pairwise path across the syntax tree reveals the structural relationship between two code tokens. Similarly, Eq.4 models the positional relationship between two words, so it is natural to modify the Eq.4 to integrate path encoding rij into self-attention module. We firstly integrate rij into the Query-Key product, that is
αij = 1√ d [(xiW Q)(xjW K + rijW K r︸ ︷︷ ︸
1
)T ], (6)
and then we integrate rij into the weighted sum of Value, that is
zi = n∑ j=1 exp (αij)∑n j′=1 exp (αij′) (xjW V + rijW V r︸ ︷︷ ︸
2
), (7)
where WKr ,W V r ∈ Rd×d are projection matrices of Key and Value for rij , respectively. For clarity, we label the modifications for Key and Value as 1 and 2, respectively.
Absolute path encoding The embed vectors of an absolute path are fed into another bi-directional GRU [Cho et al., 2014], that is
ai = GRUa(Pathxi→root), (8)
where Pathxi→root is the nonterminals vectors sequence of the absolute path from tokens xi to the syntax tree root, and ai ∈ Rd is the final vector representation of this path. The absolute path for each source code token can be viewed as a directed edge connecting the terminal with the tree root, which means the absolute path can calibrate the coordinates for each terminal in the special structural space that assigns the tree root as the origin. So it is natural to modify the absolute positional encoding in Eq.3 to integrate path vectors into the self-attention module, that is
αij = 1√ d [(xiW Q)(xjW K)T + (aiW Q a )(ajW K a ) T ], (9)
where WQa ,W K a ∈ Rd×d are projection matrices of Query and Key for absolute path representation, respectively.2
We think it is meaningful to investigate the feature interaction between absolute and relative paths and focus on both of them instead of overlooking either. So we further combine the relative and absolute path encoding shown in previous paragraphs and merge Eq.6 and Eq.9 as
αij = 1√ d [(xiW Q)(xjW K + rijW K r︸ ︷︷ ︸
1
)T + (aiW Q a )(ajW K a ) T ]. (10)
In the paragraphs above, we propose several modifications to integrate paths into self-attention. To sum up, we combine them and propose two versions of models. We firstly name the method of only relative path encoding as TPTRANS (Eq.6 + Eq.7). Then, to analyze the relationship between the relative and absolute path encoding further, we name the method of combining both of them as TPTRANS-α (Eq.10 + Eq.7).
2We tried the method adding path representation with token directly like [Kim et al., 2020] but no improvement, which we think is similar to the conclusion of untying correlations in TUPE. Besides, we tried to keep the scaling term √ 2d in Eq.3 but no improvement.

heading: 4 Experiment setup

text:
Code summarization is one of the most popular tasks in code representation learning. Given a function body, this task is to predict the function name. As observed by [Allamanis et al., 2016, Alon et al., 2019], code summarization is a good benchmark for code representation learning, as the method body typically forms complete logical units and the method name tends to be descriptive and precise. In this task, our approaches jointly learn code contexts and path representation in AST and predict the target method name as a sequence of sub-tokens. We measure precision, recall, and F1 score over target sequences in the case insensitive like [Alon et al., 2018, Zügner et al., 2021].
Datasets To show the effectiveness of our approaches across different source code languages, we experiment on four datasets introduced in the CodeSearchNet (CSN) Challenge [Husain et al., 2019]: Python, Javascript, Go, and Ruby. The datasets from CodeSearchNet have been carefully deduplicated to avoid data leakage from training sets. See Table 1 for a summary of the datasets.
Preprocessing We produce the AST for each method using the open-source parser Tree-Sitter,3 in which all the code tokens are natively mapped as terminals (see Fig1). We adopt the way to split sub-tokens following [Alon et al., 2018, Zügner et al., 2021], in which each code token is split into sub-tokens respective to code naming conventions, e.g., setConnectionsPerServer is split into [set, connections, per, server]. We limit the vocabulary of sub-tokens with at least 100 occurrences in training sets, remove all punctuations and restrict the max code length 512. We remove all anonymous Javascript functions following [Zügner et al., 2021]. We assign different ids for nonterminals and do not split the literal of them. We set the max path length 32 covering almost all paths and make padding for short paths. For paths longer than 32, we sample nodes with equal intervals to maintain max length instead of truncating the last nodes since we assume the nodes at both left and right ends of the path are equally important.
Parameter sharing In our implementation, the embedding matrix for nonterminals is shared for both absolute and relative paths since both kinds of paths are made up of nonterminals within the same syntax tree. We do not share the path encoder GRU for two different kinds of paths. Since the number of nonterminal types is far less than the size of source code vocabulary, we set different embedding matrices for nonterminals and source tokens, respectively. For efficiency, we share in different layers and heads for the projection matrices WKr ,W V r in relative path encoding and W Q a ,W K a in absolute path encoding. We also share the path representation across different layers and heads, which means we pre-encode paths only once by GRU before feeding features into Transformer.
Complexity optimization The naive way for relative path encoding needs to encode all the pairwise paths for each code snippet and to compose the relative path representation matrix R ∈ Rn×n×d, where n is code length and d is the dimension of hidden representations. This approach costs O(n2) complexity since the number of pairwise paths encoded is n(n−1)2 . To optimize it, we notice that much repetition exists in all relative paths of a sample.4 That means we only need to encode all the unique pairwise paths for each sample. We record the mapping between token pairs and unique paths as matrix M ∈ Rn×n. After that, we compute the Query-Key product between tokens and unique paths, and gather the product scores using matrix M . After calculating attention distribution, we
3https://github.com/tree-sitter/ 4For example, if we parse a function that has three formal arguments to a syntax tree, we could see that the
pairwise relative path between any two arguments will be the same.
scatter attention probability to each unique path and weighted sum using matrix M again. This new approach avoids encoding O(n2) paths and reduces it into O(l) where l is the number of unique paths and is far less than n2. It also reduces matrix R from O(n2d) to O(ld). In our experiment, we set the max l 512 for the relative path, which covers most of the samples. Besides, the complexity of absolute path encoding is original O(n), but we also reduce it in the same way and set the max l 256 for it. We refer readers to the appendix for further details.
Normalization We integrate vector sequence of path nodes as the inductive bias for the attention module of Transformer, which is not just trivial technically. Firstly, directly integrating standard RNN into the deep backbone of Transformer most likely results in gradient exploding or vanishing. Secondly, the attention procedure of Transformer is similarity calculation between vectors. The Q, K, and V of the attention module are normalized while standard RNN does not limit vector norm, leading to size instability for dot-product attention between word and path representation. To overcome these issues, we first replace the standard GRU as a layer-normalized GRU [Ba et al., 2016], getting more stable parameters updates. After getting the final state of GRU, we feed it into another normalization layer to keep it the same vector norm with Q, K, and V . The normalization is much efficient shown in the following ablation result.
Hyperparameters We denote the number of layers for encoder and decoder as LE and LD, the hidden size as D, the feed-forward dimension as DFF , and the number of heads as H . We primarily report results on two model size for both TPTrans and TPTrans-α: base (LD=1, DFF =2048) and large (LD=3, DFF =4096). For base and large settings, we all set LE=3, D=1024, and H=8. We add a pointer network [Vinyals et al., 2015] to the decoder as same as our baselines. The base model setting keeps the same size as other baselines for fair comparisons. The embedding dimension of the word and path node are 512 and 64, and we apply a linear layer to project word embedding to the hidden size of Transformer. We use one layer Bi-GRU, set the hidden size to 64, and concat the final states of forward and backward as output with the size of the single head dimension. We set the batch size and dropout rate to 128 and 0.2 and employ label smoothing of 0.1. All models and baselines are trained from random initial parameters. As the optimizer, we use Adam [Kingma and Ba, 2014] with a learning rate and weight decay of 1e−4. In our experiment, we use 4 Tesla V100 GPUs for training.
Baselines Combining two views of context and structure for source code has been widely discussed in Code Intelligence. The context mainly refers to representing code as a sequence of text, and the structure refers to extracting knowledge from code syntax. Here we mainly compare with Code2seq [Alon et al., 2018], Graph Relational Embedding Attention Transformer (GREAT) [Hellendoorn et al., 2019], XLNet [Yang et al., 2019] and Code Transformer [Zügner et al., 2021]. Unlike the Transformer-based baselines, Code2seq uses only pairwise path information in AST to represent code snippets, which means it mainly focuses on code structure within the syntax tree and lacks context modelling. On the contrary, our proposed models and other baselines are designed based on Transformer naturally focused on context. Specifically, GREAT uses the relative positional encoding present in [Shaw et al., 2018] to bias the attention via manually designed edges, and Code Transformer is a recent Transformer model based XLNet, incorporating multiple relations to learning both structure and context jointly. XLNet is designed initially for huge-corpus natural language pretraining, but its novel two-stream self-attention also shows powerful performance in code representation learning.
Essentially, the difference between our model and baselines GREAT and Code Transformer is encoding path vs encoding edges in GREAT and distances in Code Transformer. The different node combinations of paths naturally contain plenty of structure information, which is overlooked by encoding distances. See Table 2 for clear comparisons between our models and baselines.

heading: 5 Results

heading: 5.1 Overall comparison

text:
The overall comparison results are shown in Table 3. We show the performance of two different settings for both TPTrans and TPTrans-α. The base setting keeps the same Transformer backbone as other Transformer-based baselines, and the large setting additionally widens the feed-forward dimension and increases the number of decoder layers. We find that our proposed models TPTrans and TPTrans-α substantially outperform all baselines on all but one language, highlighting the effectiveness of learning path information. The only exception is Javascript, while our models still outperform most of the baselines on this dataset. Besides, we find that the large setting further improves the performance on most datasets.
It is worth mentioning that Code Transformer is designed based on XLNet, which natively equips the relative positional encoding of context and is effective for modelling long sequences. In this work, we mainly focus on studying the effectiveness of encoding path, so we use the original Transformer as our backbone. Even so, we still obtain advanced results on almost all datasets. We also find that compared to the improvement of Code Transformer over XLNet, our models get a more remarkable rise over Code Transformer.
To investigate the exception of Javascript dataset, we focus on TPTrans and count the mean and variance of all relative paths length across four datasets. We also count the total numbers of unique relative paths in training datasets and present them in Table 4.5 We firstly find that the path length of Javascript is not only longer but also has a significantly larger variance than other datasets. We then find that the unique paths num of Javascript is much large and does not match its dataset size well, which likely indicates much noise exists in the training dataset. So we assume that due to such property of Javascript, it is harder to extract useful structural information from path than others. Accordingly, our models achieve more significant performance on Ruby dataset, and the Ruby dataset has smaller mean length and variance than others and also keeps an ideal num of unique paths. As for the outlier of Javascript, one possible reasoning line is the inherent property of Javascript language results in complicated syntax. We also suspect that it is due to the internal design of language parsers, and perhaps a more suitable parser might alleviate this problem.
Lastly, we find that the overall performance of TPTrans and TPTrans-α does not differ so much, although the latter seems to obtain additional information about the absolute path. We investigate the effectiveness of both paths and analyze the relationship between them further in the following.
5For Javascript, we present the sample num of non-anonymous functions.

heading: 5.2 Empirical study

text:
In this work, we propose two versions of models, namely TPTrans and TPTrans-α. The first one encodes the pairwise relative paths in the syntax tree, while the last one additionally encodes the absolute path. Here we explore the feature interaction of relative and absolute path encodings and mainly experiment on the large model setting.
Relative path encoding For TPTrans (Eq.6 + Eq.7), We first remove the relative path encoding in the weighted sum of Value (2 in Eq.7) and then further remove the relative path in Query-Key product (1 in Eq.6), showing in the top half of Table 5. We find that with the elimination of relative path encoding gradually, the performance of TPTrans decreases obviously. After removing all submodules of relative path encoding (1 in Eq.6 and 2 in Eq.7), TPTrans totally degenerates to vanilla Transformer and performs the worst. These results highlight the usefulness of relative path encoding, both in Query-Key product (Eq.6) and the weighted sum of Value (Eq.7).
Absolute path encoding TPTrans-α extra adds absolute path encoding on Query-Key product of attention score (Eq.10) compared to TPTrans. We remove all the relative path encoding of TPTrans-α and compare it to vanilla Transformer shown in the last rows of both top and bottom halves of Table 5, respectively. We find that absolute path encoding significantly improves vanilla Transformer, highlighting its effectiveness.
Relationship between them We analyze the interaction between these two kinds of paths further. In the bottom half of Table 5, we firstly find that with the elimination of two submodules of relative path encoding gradually (1 in Eq.10 and 2 in Eq.7), the performance of TPTrans-α decreases at the same time. This phenomenon is consistent with the ablation study for relative path encoding shown in the previous paragraph. We then compare the top and bottom halves of Table 5 row by row from down to up and find that with the introduction of relative path gradually, the performance gap between TPTrans and TPTrans-α correspondingly decreases, which means the improvement brought by absolute path becomes smaller at the same time. After introducing full of the relative path encoding, the gain from absolute path encoding is almost weak. 6 These results show that these two kinds of paths are not orthogonal and much feature overlap indeed exists between them. Actually, if two sub-tokens are too far away from each other in the syntax tree, the relative path between them then spans the root of AST, in which the relative and absolute path encodings play the same role for modelling structure. The experimental results also show that it is preferable to encode the relative path than the absolute path.

heading: 5.3 Ablation study

text:
We explore the roles of each part of our approaches. We choose TPTrans with the base setting and mainly experiment on Python and Ruby datasets.
6We also tried another method adding absolute path representation to word embedding directly in [Kim et al., 2020] as mentioned before, but the overall conclusion remains all the same.
Path vs. distance To verify the benefit of learning from paths than distances presented in [Zügner et al., 2021], we try to convert all nodes in paths to the same one before feeding paths into TPTrans and present in Table 6. After such special preprocessing, what TPTrans learns completely degenerates into hops across the syntax tree between two terminals. We find that the performance of our proposed model shows noticeable declines after information ablation. It indicates TPTrans indeed learns from rich node combinations within paths, and it is better to model path than distance.
Normalization As mentioned before, embedding vector sequences into the attention module of Transformer is not just trivial, so we carefully design the normalization method and present the ablation results in Table 7. These results prove that the normalization method is crucial for embedding vector sequence into attention, which is also helpful to guide the design of neural networks.
Deeper model architecture We re-design a deeper model architecture for TPTrans and vanilla Transformer and set LE=LD=6, D=512, DFF =2048 and H=8. For TPTrans, the input and output sizes of GRU are both 64. We present the results in Table 8 and conclude that our approaches of integrating path encodings still work for the deep model setting.

heading: 6 Conclusion

text: In this paper, we investigate the interaction between the absolute and relative path encodings by integrating them into the unified Transformer framework, and confirm that feature overlap exists between these two kinds of paths. The detailed empirical study for path encodings also leads to our novel code representation model TPTrans. Extensive experiments and ablation study demonstrate the effectiveness of our approaches on code summarization across four different programming languages.

heading: 7 Acknowledgements

text: We thank all reviewers for their constructive comments. This research is supported by the National Key R&D Program of China under Grant No. 2020AAA0109400, and the National Natural Science Foundation of China under Grant Nos. 62072007, 61832009, 61620106007.

title: Integrating Tree Path in Transformer for Code Representation