AST-based Neural Network (ASTNN)
J. Zhang, X. Wang, H. Zhang, H. Sun, K. Wang and X. Liu, “A Novel Neural Source Code Representation Based on Abstract Syntax Tree,” 2019 IEEE/ACM 41st International Conference on Software Engineering (ICSE), 2019, pp. 783-794, doi: 10.1109/ICSE.2019.00086.
背景
-
过往的基于抽象语法树(AST)的神经模型可以较好地表示代码,也广泛用于代码分类、克隆检测等。但这些神经模型生成的AST规模巨大,同时常存在严重的长期依赖性问题。
-
简单提一提长期依赖性问题:假设一个简单RNN的循环联系是:
a < t > = W T a < t − 1 > a^{<t>}=W^Ta^{<t-1>} a<t>=WTa<t−1>
⟹ a < t > = ( W t ) T a < 0 > \Longrightarrow a^{<t>}=(W^t)^Ta^{<0>} ⟹a<t>=(Wt)Ta<0> -
其中 W W W 常满足: W = Q Λ Q T W=Q\Lambda Q^T W=QΛQT, Q Q Q 为正交矩阵
-
则有:
a < t > = Q T Λ t Q a < 0 > a^{<t>}=Q^T\Lambda ^tQa^{<0>} a<t>=QTΛtQa<0> -
如此一来,随着 t t t 递增,幅值小于1的衰减为0,幅值大于1的急剧增大。缺少非线性激活函数的RNN问题更为严重。
-
对于处理AST这类树形结构,最近较新、应用较广的有三种代表性模型:递归神经网络(RvNN),基于树的卷积神经网络(Tree-based CNN, TBCNN),以及基于树的长短期记忆模型(Tree-LSTM)。
方法
- 首先,论文中提到,当Java或C语言产生的AST达到约10000个节点、深度约为100的时候,滑动窗口已经会产生较明显的长期信息丢失;同时以上三种方法都将AST转化成完全二叉树处理,这破坏了源代码的原始语法结构,进一步使得AST更加庞大。
语句树(ST-tree)的拆分
- 本文提出的ASTNN不会一次性地考虑整棵AST,会选择先将AST拆分成若干较小的语句树,再将这些小型语句树编码成向量。
-
如何将语句树编码是本文的重要内容,对于一棵语句树t,首先可以通过以下公式来获取一个非叶子节点n的词汇向量表示:
v n = W e T x n v_n=W_e^Tx_n vn=WeTxn -
其中 x n x_n xn 是节点 n 的 one-hot 表示, v n v_n vn 是其embedding, 嵌入参数 W e ∈ R ∣ V ∣ ∗ d W_e\in R^{|V|*d} We∈R∣V∣∗d (词汇量V,嵌入维度d) 是预训练好的。
-
随后我们可以如此获得节点n的向量表示:
h = σ ( W n T v n + Σ i = 1 C h i + b n ) h=\sigma (W_n^Tv_n+\Sigma_{i=1}^{C}{h_i}+b_n) h=σ(WnTvn+Σi=1Chi+bn)
-
其中 W n ∈ R d ∗ k W_n\in R^{d*k} Wn∈Rd∗k 是编码维度为 k 的权重矩阵, b n b_n bn 是偏项,C 是节点 n 的子节点数目,h 表示隐藏状态, σ \sigma σ 是激活函数,通常可用tanh或者恒等式。
-
经过最大值池化,最终一个语句树的向量可表示为:(N 是该语句树的节点个数)
e t = [ m a x ( h i 1 ) , … , m a x ( h i k ) ] , i = 1 , … , N e_t=[max(h_{i1}),…,max(h_{ik})], i=1, …, N et=[max(hi1),…,max(hik)],i=1,…,N
动态批处理(dynamic batch)
- 有意思的是,本文针对大型数据集的训练效率,设计了一套批处理算法来进行一定程度的并行计算。原理比较简单,在递归时分层处理,针对每层节点的子节点归类处理。
表示语句序列
- 该部分主要用到了门控循环单元(Gated recurrent unit, GRU)
r t = σ ( W r e t + U r h t − 1 + b r ) z t = σ ( W z e t + U z h t − 1 + b z ) h ~ t = t a n h ( W h e t + r t ⊙ ( U h h t − 1 ) + b h ) h t = ( 1 − z t ) ⊙ h t − 1 + z t ⊙ h ~ t (4) \begin{aligned} &r_{t}=\sigma(W_{r}e_{t}+U_{r}h_{t-1}+b_{r})\\ &z_{t}=\sigma(W_{z}e_{t}+U_{z}h_{t-1}+b_{z})\\ &\tilde{h}_{t}=tanh(W_{h}e_{t}+r_{t}\odot(U_{h}h_{t-1})+b_{h})\\ &h_{t}=(1-z_{t})\odot h_{t-1}+z_{t}\odot \tilde{h}_{t}\\ \tag{4} \end{aligned} rt=σ(Wret+Urht−1+br)zt=σ(Wzet+Uzht−1+bz)h~t=tanh(Whet+rt⊙(Uhht−1)+bh)ht=(1−zt)⊙ht−1+zt⊙h~t(4)
-
其中, r t r_t rt 是用于控制先前状态影响的 重置门, z t z_t zt 是用于组合过去信息和新信息的 更新门, W _ , U _ W\_\ ,U\_ W_ ,U_ 都是权重矩阵, b _ b\_ b_ 都是偏差项, h ~ t \tilde h_t h~t 是候选状态,用于与前状态 h t − 1 h_{t-1} ht−1 线性插值得到当前状态 h t h_t ht 。
-
也很有意思的是,为了进一步增强递归层捕获依赖信息的能力,本文采用了双向 GRU 技术,即:将两个方向得到的 h t h_t ht 组合得到一个隐藏状态。
h t → = G R U ⟶ ( e t ) , t ∈ [ 1 , T ] h t ← = G R U ⟵ ( e t ) , t ∈ [ T , 1 ] h t = [ h t → , h t ← ] , t ∈ [ 1 , T ] (5) \begin{aligned} &\overset{\rightarrow}{h_{t}}=\overset{\longrightarrow}{GRU}(e_{t}), t\in[1,\ T]\\ &\overset{\leftarrow}{h_{t}}=\overset{\longleftarrow}{GRU}(e_{t}), t\in[T,\ 1]\\ &h_{t}=[\overset{\rightarrow}{h_{t}},\ \overset{\leftarrow}{h_{t}}], t\in[1,\ T]\\ \tag{5} \end{aligned} ht→=GRU⟶(et),t∈[1, T]ht←=GRU⟵(et),t∈[T, 1]ht=[ht→, ht←],t∈[1, T](5)
- 最后,这些状态都会进入池化层采样。为了区分不同语句的重要性,本文使用了最大值池化。