[論文レビュー] TreeBERT: A Tree-Based Pre-Trained Model for Programming Language

Xue Jiang, Zhuoran Zheng+3 more|arXiv (Cornell University)|Jul 27, 2021

Software Engineering Research被引用数 1

ひとこと要約

TreeBERT は、抽象構文木（AST）を構成経路として表現し、ノードの位置埋め込みを用いることで、コード内の AST を木構造に基づいて活用する事前学習モデルである。木マスク言語モデル化（TMLM）とノード順序予測（NOP）を用いてプログラミング言語生成を向上させ、コード要約およびドキュメンテーションタスクにおいて、未学習のプログラミング言語へのゼロショット転送を含め、既存のモデルを上回る性能を発揮する。

ABSTRACT

Source code can be parsed into the abstract syntax tree (AST) based on defined syntax rules. However, in pre-training, little work has considered the incorporation of tree structure into the learning process. In this paper, we present TreeBERT, a tree-based pre-trained model for improving programming language-oriented generation tasks. To utilize tree structure, TreeBERT represents the AST corresponding to the code as a set of composition paths and introduces node position embedding. The model is trained by tree masked language modeling (TMLM) and node order prediction (NOP) with a hybrid objective. TMLM uses a novel masking strategy designed according to the tree's characteristics to help the model understand the AST and infer the missing semantics of the AST. With NOP, TreeBERT further extracts the syntactical structure by learning the order constraints of nodes in AST. We pre-trained TreeBERT on datasets covering multiple programming languages. On code summarization and code documentation tasks, TreeBERT outperforms other pre-trained models and state-of-the-art models designed for these tasks. Furthermore, TreeBERT performs well when transferred to the pre-trained unseen programming language.

研究の動機と目的

• 事前学習に木構造の AST を組み込むことで、プログラミング言語生成を向上させること。
• 既存のコード用事前学習モデルにおける木構造の活用不足に取り組むこと。
• コードの AST における意味的・構文的構造を同時に捉えるハイブリッド事前学習目的を考案すること。
• 強固な構造的表現を通じて、未学習のプログラミング言語への効果的なゼロショット転送を可能にすること。
• 木に配慮した表現学習を通じて、コード要約およびドキュメンテーション生成を向上させること。

提案手法

• 木構造の保持と木に配慮したモデリングを可能にするために、コードの AST を構成経路として表現すること。
• AST 構造内でのノードの相対的位置を符号化するために、ノード位置埋め込みを導入すること。
• 欠落した AST ノードを予測するための木特有のマスキング戦略を用いた木マスク言語モデル化（TMLM）目的の設計。
• AST 内のノードの正しい順序を予測することで構文的制約を学習するため、ノード順序予測（NOP）を提案すること。
• TMLM と NOP を統合したハイブリッド事前学習目的を設計し、意味的および構造的理解を同時に最適化すること。
• 多言語コードデータセット上で TreeBERT を事前学習することで、プログラミング言語間での一般化を向上させること。

実験結果

リサーチクエスチョン

• RQ1 明示的に AST 構造をモデル化する事前学習モデルは、要約やドキュメンテーションといったコード生成タスクを向上させることができるか？
• RQ2 木特有のマスキングを用いた木マスク言語モデル化（TMLM）は、AST の意味をどれほど効果的に捉えられるか？
• RQ3 ノード順序予測（NOP）は、コード内の構文的構造理解をどの程度向上させるか？
• RQ4 TreeBERT は、未学習のプログラミング言語へのゼロショット転送に一般化できるか？
• RQ5 TreeBERT は、コード生成ベンチマークにおいて、既存の事前学習モデルや最先端手法と比較してどの程度優れているか？

主な発見

• TreeBERT は、複数のプログラミング言語にわたるコード要約およびドキュメンテーション生成タスクにおいて、既存の事前学習モデルを上回る性能を発揮する。
• TMLM と NOP のハイブリッド事前学習目的は、単独で使用する場合よりも優れた性能をもたらす。
• TreeBERT は未学習のプログラミング言語においても強力なゼロショット転送性能を達成し、構造的一般化の優れた能力を示している。
• TMLM における木特有のマスキング戦略は、AST 内の欠落した意味的コンテンツの推論能力を顕著に向上させる。
• ノード順序予測（NOP）は構文的制約を効果的に捉えており、モデルのコードに対する構造的理解を強化している。
• TreeBERT は、コード生成のベンチマークデータセットにおいて最先端の結果を達成しており、木に基づく事前学習の有効性を裏付けている。