[論文レビュー] Learning Program Representations with a Tree-Structured Transformer

Wenhan Wang, Kechi Zhang+5 more|arXiv (Cornell University)|Aug 18, 2022

Software Engineering Research被引用数 6

ひとこと要約

この論文では、親子ノードおよび兄弟ノード間の依存関係を捉えるためにマルチヘッドアテンションを活用し、双方向（下向きおよび上向き）の情報伝播を組み合わせることで、プログラム表現を学習する新しいツリー構造のトランスフォーマーモデルであるTree-Transformerを提案する。この手法は、長距離依存関係を効果的にモデル化し、豊富なノードレベル表現を学習することで、プログラム分類、バグ局在、型推論のタスクにおいて、従来のツリー基盤およびグラフ基盤の手法を顕著に上回る性能を発揮する。

ABSTRACT

Learning vector representations for programs is a critical step in applying deep learning techniques for program understanding tasks. Various neural network models are proposed to learn from tree-structured program representations, e.g., abstract syntax tree (AST) and concrete syntax tree (CST). However, most neural architectures either fail to capture long-range dependencies which are ubiquitous in programs, or cannot learn effective representations for syntax tree nodes, making them incapable of performing the node-level prediction tasks, e.g., bug localization. In this paper, we propose Tree-Transformer, a novel recursive tree-structured neural network to learn the vector representations for source codes. We propose a multi-head attention mechanism to model the dependency between siblings and parent-children node pairs. Moreover, we propose a bi-directional propagation strategy to allow node information passing in two directions, bottom-up and top-down along trees. In this way, Tree-Transformer can learn the information of the node features as well as the global contextual information. The extensive experimental results show that our Tree-Transformer significantly outperforms the existing tree-based and graph-based program representation learning approaches in both the tree-level and node-level prediction tasks.

研究の動機と目的

• コードの表現学習において、長距離依存関係を捉えることや効果的なノードレベル表現を学習する点で、従来のツリー基盤およびグラフ基盤のモデルに見られる限界を克服すること。
• リーフノードおよび低レベルノードの文脈学習を妨げる、従来の再帰的ニューラルネットワークにおける一方向の情報フローを克服すること。
• 特に無順序性を考慮しないモデル（例：GNN）が無視しがちな、プログラムの順序付き構造、特に兄弟ノードの順序を保持・活用すること。
• 複雑な前処理やグラフ拡張を必要とせず、元の構文木構造を維持することで、オーバーヘッドを低減し、構造の正確性を保つこと。
• 階層的かつ順序付きのプログラム構造を効果的にモデル化するトランスフォーマーベースのアーキテクチャを開発し、プログラム表現学習の向上を図ること。

提案手法

• Tree-Transformerは、ツリー構造に特化したマルチヘッド自己アテンション機構を採用し、親子ノードおよび兄弟ノードペア間の依存関係を明示的にモデル化する。
• 下向きおよび上向きの両方向情報伝播を実現するため、2つの並列なトランスフォーマーユニット（下向きユニット：リーフからルートへ、上向きユニット：ルートから子孫へ）を導入する。
• ノードの兄弟順序を捉えるために位置エンコーディングを統合し、木構造内の異なる構造的位置を区別できるようにする。
• アテンションベースの統合を用いて子ノード表現の再帰的集約によりノード表現を学習し、木全体にわたり文脈情報が流れ込むようにする。
• 再帰的かつツリー構造的なアーキテクチャを採用しており、抽象構文木（AST）から直接ベクトル表現をエンドツーエンドで学習可能である。
• プログラム分類、バグ局在、型推論などの下流タスクで標準的な教師あり学習目的関数を用いて学習を行う。

Figure 1: The overall pipeline of Tree-Transformer.

実験結果

リサーチクエスチョン

• RQ1 トランスフォーマーベースのアーキテクチャは、特に離れたノード間においても、プログラム構文木における長距離依存関係を効果的にモデル化できるか？
• RQ2 下向きおよび上向きの双方向情報伝播は、一方向手法と比較してノードレベル表現学習を向上させられるか？
• RQ3 位置エンコーディングによる兄弟順序の統合は、コードにおける構文的構造の区別能力を向上させるか？
• RQ4 下流のプログラム理解タスクにおいて、最先端のツリー基盤およびグラフ基盤のモデルと比較してTree-Transformerはどのように性能を発揮するか？
• RQ5 複雑なグラフ構築や構造的前処理を必要とせず、性能を維持できるか？

主な発見

• Tree-Transformerは、複数のベンチマークでプログラム分類タスクにおいて最先端の性能を達成し、従来のツリー基盤およびグラフ基盤のモデルを上回る。
• 誤った演算子局在や修復タスクといったノードレベル予測タスクにおいても顕著な向上を示し、ノードレベルでの表現学習の優位性を裏付けている。
• 双方向伝播機構により、遠く離れたコード要素間の文脈的推論を要するタスクで改善が見られ、長距離依存関係の効果的モデル化が実証された。
• 位置エンコーディングによる兄弟順序の統合は、構文的順序に敏感なタスクで明確な性能向上をもたらした。
• 追加のグラフ構築や前処理を必要とせず、元のAST構造を維持し、計算オーバーヘッドを低減する点で優れた性能を発揮した。
• アブレーションスタディの結果、双方向伝播およびマルチヘッドアテンションの両コンponentがモデル性能に不可欠であることが確認された。

Figure 2: An example of existing tree-structured and graph neural networks for learning on abstract syntax trees.