AST-Tran：使用树形结构的代码摘要模型

AST-Trans: Code Summarization with Efficient Tree-Structured Attention

原文链接：https://ieeexplore.ieee.org/document/9794079

摘要

代码摘要旨在为源代码生成简短的自然语言描述，最先进的方法遵循基于变压器的编码器-解码器架构，抽象语法树（AST）被广泛用于编码结构信息。然而，AST一般而言很大，由于每个节点都需要对AST中的所有其他节点计算self-attention，因此这也会导致大量的计算开销。

现有的方法忽略了尺寸约束，并简单地将整个线性化AST馈送到编码器中，这样一个简单的过程很难从过长的输入序列中提取真正有用的依赖关系。

本文中提出了AST-Trans，它利用了AST中的两种类型的节点关系：祖先关系和兄弟关系，为相关节点动态分配权重，并基于这两种关系排除无关节点。

介绍

背景：代码摘要能够帮助程序员快速消化代码，而不用自己遍历代码，然而，维护高质量的代码摘要实际上需要昂贵的人工。在许多项目中，这些总结往往不匹配、缺失或过时，从而减缓了开发进度，自动代码摘要可以大大节省开发人员的时间。

AST通常通过不同的算法进行线性化，如先序遍历、基于结构的遍历、路径分解等，然而，线性化的AST由于包含额外的结构化信息，比其相应的源代码序列长得多，这使得模型极难从过长的输入序列中准确地检测有用的依赖关系，且来大量的计算开销。

假设受其祖先节点和兄弟节点的影响最大，如上图所示，祖先节点代表层次关系，兄弟节点代表块内关系。捕获这两个信息就足以生成摘要，而不需要在所有节点之间建立self-attention。

普通线性AST树的时间复杂度是 $O(n^2)$，本文提出的复杂度为 $O(n)$。

AST-TRAN

这节将利用AST树生成祖先关系矩阵和兄弟关系矩阵，取代原有的自我注意力，并动态排除无关节点降低计算成本。

AST 线性化

为了对树形AST进行编码，首先需要使用线性化方法将其转换为序列。当前工程中使用的三种最具代表性的线性化方法：

• 先序遍历（POT）：这种方法会损失原有的AST信息，不能利用先序遍历还原出原来的AST树。
• 结构遍历（SBT）：添加了额外的括号来表征父子关系，这样可以利用这些有序对还原出AST，但是线性化序列大小增倍。
• 路径分解（PD）：通过连接两个随机叶节点之间的路径来表示AST。路径的总数对于计算来说可能太大，因此需要随机采样。

下表展示了这几种遍历方式的结果

关系矩阵

定义了两种我们关心的节点之间的关系：祖先关系A和兄弟关系S，维度都是 $N\times N$，N为节点个数。假定在线性AST树（先序遍历）中的第i个节点为 $n_i$，$A_{ij}$ 表示两节点 $n_i,n_j$之间最短路径的长度，$S_{ij}$表示两兄弟节点之间的水平距离，如果不满足这些关系，用Inf无穷大表示。

更准确的表示形式如下，其中SPD和SID分别表示到根节点的最短路径和兄弟节点之间的水平距离。
$$\begin{array}{rl}& A_{ij}=\{\begin{array}{rl}\mathbf{SPD}(i,j)& \text{ if }|\mathbf{SPD}(i,j)|\le P\\ \infty & \text{ otherwise }\end{array}\\ & S_{ij}=\{\begin{array}{rl}\mathbf{SID}(i,j)& \text{ if }|\mathbf{SID}(i,j)|\le P\\ \infty & \text{ otherwise }\end{array}\end{array}(1)$$
P是预先定义的阈值，如果距离大于这个值则被忽略，这假设了一定距离之外的节点，它们的关系是无用的，可以降低计算复杂度。

下图展示了这两个矩阵的例子，距离是有方向的

树形结构注意力

self-attention：对于标准的自注意力，假定输入为$\mathbf{x}=(x_1,\cdots ,x_n)$， 其中 $x_i\in {\mathbb{R}}^d$，是嵌入向量，输出矩阵 $\mathbf{o}=(o_1,\cdots ,o_n),o_i\in {\mathbb{R}}^d$，公式如下
$$\begin{array}{rl}{\mathbf{\alpha }}_{ij}& =\frac{Q\left(x_i\right)\mathbf{K}{\left(x_j\right)}^{\top }}{\sqrt{d}}\\ o_i& =\sum _{j=1}^n\sigma \left({\mathbf{\alpha }}_{ij}\right)\mathbf{V}\left(x_j\right)\end{array}(2)$$
其中 Q、K表示query函数和key函数，都是 ${\mathbb{R}}^d\to {\mathbb{R}}^m$，V 时value函数，${\mathbb{R}}^d\to {\mathbb{R}}^d$，$\sigma$ 是得分函数（softmax）。

Relative position embedding：上面的等式没有任何位置信息的，而相对位置嵌入在代码摘要任务中显示出更有效，相对位置 $\delta (i,j)$反应了节点 $n_i,n_j$ 之间的相对距离，假设P是最大的相对距离，则 $\delta (i,j)$ 有如下定义
$$\delta (i,j)=\{\begin{array}{rll}0& \text{ for }& i-j\le -P\\ 2P& \text{ for }& i-j\ge P\\ i-j+P& \text{ others. }& \end{array}(3)$$
通过这种方式，我们可以将每个相对距离映射到嵌入区域，可以在等式2的顶部添加相对位置嵌入，以表明成对距离。

Disentangled Attention（分散注意力）：分离注意力使用相对位置嵌入作为自我注意过程中的偏差。每个单词都使用两个向量来表示，这两个向量对其内容和相对位置进行了编码。然后将注意力计算分为三个部分：内容对内容、内容对位置和位置对内容，定义如下：
$${\tilde{\alpha }}_{i,j}=\underset{\text{content-to-content }}{\underbrace{Q\left(x_i\right)K{\left(x_j\right)}^{\top }}}+\underset{\text{content-to-position }}{\underbrace{Q\left(x_i\right)K_{\delta (i,j)}^P}}+\underset{\text{position-to-content }}{\underbrace{Q_{\delta (j,i)}^{P}K{\left(x_j\right)}^{\top }}}(4)$$
其中 $Q^P,K^P\in {\mathbb{R}}^{(2P+1)\times m}$ 表示相对位置的query矩阵和key矩阵，$K_{\delta (i,j)}^P$表示 $K^P$ 的 $\delta (i,j)$行，$Q_{\delta (i,j)}^P$表示 $Q^P$ 的 $\delta (i,j)$行，后面两个乘积，即content-to-position和position-to-content用于测量单词对之间的相对位置。时间复杂度为 $O(2Pm)$。

Attention with Tree-Structured Relationships：文章提出的距离取代了相对距离 $\delta (i,j)$，而用的 ${\delta }_R(i,j)$，其中 R 代表祖先矩阵A 或者 兄弟矩阵S，其定义如下
$${\delta }_R(i,j)=\{\begin{array}{rll}{R}_{ij}+P+1& \text{ if }& R_{ij}\in [-P,P]\\ 0& \text{ if }& R_{ij}=\infty \end{array}(5)$$
由于有两种关系，我们认为每个head中只有一种关系，这两种关系的信息将通过多头注意力机制合并在一起，最终的输出为：
o ~ i = ∑ j j ∈ { j ∣ δ R ( i , j ) > 0 } σ ( α ~ i , j 3 d ) ( V ( x j ) + V R i j P ) ( 6 ) \tilde{o}_i=\sum_j^{j \in\left\{j \mid \delta_R(i, j)>0\right\}} \sigma\left(\frac{\tilde{\alpha}_{i, j}}{\sqrt{3 d}}\right)\left(V\left(x_j\right)+V_{R_{i j}}^P\right)\qquad\qquad(6) o~i​=j∑j∈{j∣δR​(i,j)>0}​σ(3d ​α~i,j​​)(V(xj​)+VRij​P​)(6)
值得一提的是只计算 ${\delta }_R(i,j)>0$的，类似于滑动窗口的思想来减少时间和空间复杂度。$\alpha$ 要除以 $\sqrt{3d}$ 是因为有三项，要除以这个因子。 $V^P$表示value相对位置的投影矩阵， $V_{R_{ij}}^P$ 表示 $V^P$ 的第 $R_{ij} $ 行。

高效实现

标准Transformers中的全注意力机制随序列长度呈二次方增长，以前有人提出应用滑动窗口将注意力限制在固定范围内，使用滑动窗口，可以将序列数据中的节点对规划为线性分布，并与矩阵划分并行计算，这个方法不适用于AST，因为相关节点的位置分布随每个树结构而变化，本节有5中替代方法。

Mask：计算完注意力得分后，遮住 ${\delta }_R(i,j)=0$ 的节点的注意力得分，这并不能减少时间复杂度

Loop：只在 ${\delta }_R(i,j)>0$ 的节点中计算注意力得分，这可以减少计算注意力得分的时间和空间复杂度

Sparse：将 ${\delta }_R$ 以稀疏矩阵 $ST({\delta }_R)$ 存放并且放入深度学习框架，在计算矩阵乘法时能够自动跳过零元素，这只能用于content-to-position 和 position-to-content

Gather with COO(GC)：content-to-content可以通过附加的聚集操作来优化，核心思想是将需要计算的查询关键字对放入一对一的对应关系中，并将它们存储为密集矩阵。COO是存储稀疏张量的常用方法，其中只有非零元素存储为元素索引值 $CO{O}_{row},CO{O}_{col}$ 和其对应的值 $CO{O}_{val}$。

用COO矩阵来聚合前面的$Q(x),K(x),Q^P,K^P$:
$$\begin{array}{c}{Q}_{\text{row }}=Q(x)\left[CO{O}_{\text{row }};:\right];K_{col}=K(x)\left[CO{O}_{col};:\right]\\ Q_{val}^P=Q^P\left[CO{O}_{val};:\right];K_{val}^P=K^P\left[CO{O}_{val};:\right]\end{array}$$
这样，计算得分的公式就变成下面这样
$${\alpha }_{\text{coo }}=Q_{\text{row }}\odot K_{\text{col }}+Q_{\text{row }}\odot K_{\text{val }}^P+Q_{\text{val }}^P\odot K_{\text{col }}$$
$\odot$ 表示点乘，${\alpha }_{coo}$ 表示 $\tilde{\alpha }$ 中的非零元素。并且 $\tilde{\alpha }[CO{O}_{row}[i];CO{O}_{col}[i]]={\alpha }_{coo}[i]$。

Gather with decomposed COO (GDC)：在进行聚合操作钱，可以先分解COO再聚集，为了减少GC中聚集的数量，使得每个子矩阵 ${\delta }_R^s$ 只包含具有相同相对距离s的节点对

如上图3所示，原始的 ${\delta }_R$ 是最右边的图，它包含有三种值 1，2，3，按照值将矩阵拆解成三个子矩阵，再对这三个子矩阵 $CO{O}^s$ 进行 GC 操作：
$$Q_{ro{w}_s}=Q(x)[CO{O}_{row}^s;:],K_{co{l}_s}=K(x)[CO{O}_{col}^s;:]$$
自注意力得分就如下计算得到：
$${\alpha }_{co{o}_s}=(Q_{ro{w}_s}+Q_s^P)\odot (K_{ro{w}_s}+K_s^P)-(Q_s^P\odot K_s^P)$$
其中 ${\alpha }_{co{o}_s}$ 表示${\delta }_R^s$ 中节点的自注意力得分，所有子矩阵的自注意力得分可以并行计算，最终的自注意得分就等于所有 ${\alpha }_{co{o}_s}$ 的和：
$${\alpha }_{coo}=\sum _{s=1}^{2P+1}{\alpha }_{co{o}_s}$$
使用分解的GC的原因有三个：

• $K^P,Q^P$ 可以重复使用，因为每一个 $Q_{ro{w}_s}，K_{ro{w}_s}$ 有相同的相对距离 s。位置嵌入 s 可以直接添加到内容content中而不需要收集操作。
• 只需要四分之一的收集操作（这个文中有证明）
• 只需要点乘，且 $Q_s^P\odot K_s^P$ 的结果可以重复利用

实验

实验设置

实验训练集是Java和Python，下图显示了大多数祖先关系距离在5以内，兄弟关系距离在10以内

预处理阶段：删除标点、分割单词，拆分的节点视为原始父节点的新子节点，例如gettabletypes作为父节点，get、table、types作为该节点的三个子节点，还会“reverse the children of the root node”当AST大小超过指定的最大大小，来防止重要信息被剪切。

评估指标

使用了语料库 BLEU、METEOR、ROUGE-L 来评估性能。

基线方法

与16中方法进行比较，根据输入类型可以分为5组：

• Code：直接将代码视为纯文本作为输入
• AST：以AST作为输入，并使用特定于树的编码器进行编码，如Tree LSTM 为训练添加代码信息和强化，另一种将AST视为图形进行编码
• AST(PD)：以路径分解的AST作为输入
• AST(SBT)：结构化AST
• AST(POT)：线性化AST

实验结果

如下表3所示，AST-Trans在三个指标都优于所有基线

结果表明，当使用了祖先关系和兄弟关系时，性能得到了提高。然而，单独使用其中一种关系已经可以取得接近的结果，并超过所有先前的基线。

随着兄弟关系和祖先关系最大距离的增大，性能也会得到提高，但是由于它们大部分距离在10和5之间，所以越往后，改善越来越小