コード生成のためのツリーベースのTransformerアーキテクチャ

導入：

コード生成は、プログラマーの生産性を大幅に向上させる可能性を秘めた重要な AI 問題です。自然言語で記述された仕様が与えられると、コード生成システムはその仕様を実行可能なプログラムに変換します。たとえば、Python プログラマーが「辞書 Dict を初期化する」という命令を与えた場合、コード ジェネレーターは自動的に「Dict={}」を生成する必要があります。

ディープラーニング技術の発展に伴い、研究者はシーケンスツーシーケンス（Seq2Seq）モデルやシーケンスツーツリー（Seq2Tree）モデルなどのさまざまなニューラルアーキテクチャをこの問題に適用してきました。特に、最先端の方法では、文法規則のシーケンスを予測することによってコードを生成します。つまり、システムは生成されたコードの部分的な抽象構文木 (AST) を保持し、特定のノードを展開するために使用される文法規則を予測します。

文法規則の分類には、主に 2 つの課題があります。最初の課題は長期依存の問題です。コード要素は別の離れた要素に依存する場合があります。たとえば、100行目の変数参照文「if len(a)」

本稿では、コード生成のための新しいニューラル アーキテクチャ TreeGen を提案します。最初の課題に対処するために、TreeGen は長期的な依存関係をキャプチャできる、最近提案された Transformer アーキテクチャを採用しています。しかし、元の Transformer アーキテクチャはプログラム用に設計されておらず、前述の 2 番目の課題であるツリー構造を活用できません。グラフベースおよびツリーベースの畳み込みニューラル ネットワークと同様に、構造情報を活用する標準的な方法は、ノードとその近傍のベクトル表現を、構造化された畳み込みサブレイヤーの出力として組み合わせることです。ただし、標準の Transformer アーキテクチャには、このような構造化された畳み込みサブレイヤーが存在せず、どこに追加すればよいのかが明確ではありません。

すべての Transformer ブロックに構造化された畳み込みサブレイヤーを追加したくなるかもしれません。私たちの中心的な推測は、ノードとその隣接ノードに対して畳み込みを実行する場合、ベクトル表現には主に元のノードの情報が含まれるはずだということです。ノードのベクトル表現がトランスフォーマーのデコーダー内のより多くのブロックによって処理されるにつれて、他のノードからのより多くの情報が徐々に混ざり合い、元の情報が失われます。したがって、構造化畳み込みサブレイヤーを、すべての Transformer デコーダー ブロックではなく、最初のいくつかの Transformer デコーダー ブロックにのみ追加します。

一般に、TreeGen アーキテクチャは 3 つの部分で構成されます: (1) 自然言語 (NL) エンコーダ: テキスト記述をエンコードします。(2) AST デコーダ (最初のいくつかの Transformer Decoder ブロック) は、構造化畳み込みサブレイヤーを使用して、以前に生成された部分コードをエンコードします。(3) デコーダ (残りの Transformer Decoder ブロック) は、クエリ (AST で展開されるノード) を最初の 2 つのエンコーダと組み合わせて、次の文法規則を予測します。

私たちは、カードゲーム「ハースストーン」の Python 実装である、確立されたベンチマーク データセットでモデルを評価しました。結果は、私たちのモデルが以前のモデルよりも 4.5 パーセントポイント大幅に優れていることを示しています。さらに、自然言語の文章をラムダ計算の論理形式に変換する 2 つの意味解析データセット (ATIS と GEO) でモデルを評価します。結果は、私たちのモデルが以前のニューラル モデルの中でそれぞれ 89.1% と 89.6% という最高の精度を持っていることを示しています。私たちの評価では、最初のいくつかの Transformer ブロックに構造化畳み込みサブレイヤーを追加すると、すべてのブロックに構造化畳み込みがある Transformer よりも大幅に優れたパフォーマンスが得られることも示されています。

私たちのモデル:

プログラミング言語の文法規則を予測してコードを生成します。図 2 は、NL エンコーダー、AST エンコーダー、デコーダーの 3 つの部分で構成されるモデルの全体図を示しています。以下のサブセクションで詳細を紹介します。

文法規則の予測:

このセクションでは、一連の文法規則の分類問題としてコード生成をモデル化する方法を紹介します。プログラムは、いくつかの文脈自由文法規則に分解され、AST に解析されます。たとえば、図 1 はコード「length=10」の Python AST を示しています。破線のボックスはターミネータで、実線のボックスは非ターミネータです。

AST ベースのコード生成は、文法規則を介して非終端記号を展開するものと見なすことができます。このプロセスは、すべてのリーフ ノードが終了するまで繰り返されます。図 1 の「1:root->Module」は文法規則の例であり、先頭の数字は規則の ID です。意図したトラバーサルを実行すると、右上隅に表示される AST を生成するルールのシーケンスを取得できます。

正式には、確率は、次の順序でコードを生成するための規則に従う確率に分解できます。ここで、ri はルール シーケンス内の i 番目のルールです。このように、私たちのタスクは、p(ri|NL input, pi)を計算するモデルをトレーニングすることです。つまり、自然言語記述と現在生成されている部分的なASTが与えられると、モデルはそのノードを拡張するルールの確率を計算します。

NL エンコーダー:

入力した説明によって、コードの動作が決まります。これは、Hearthstone データセット内の半構造化記述、または ATIS および GEO セマンティック解析データセット内の自然言語になります。

入力を説明するために、まず入力を n1、n2、...、nL とラベル付けします。ここで、L は入力の長さを表します。次に、各 ni は文字 c1(ni)、c2(ni)、...、cS(ni) に分割されます。ここで、S は ni 内の文字数です。埋め込みにより、すべてのトークンと文字は数値ベクトル n1、n2、...、nL および c1(ni)、c2(ni)、...、cS(ni) として表されます。

以下を示すテキストを入力してください:

文字の埋め込み。類似した単語には類似した文字が含まれることがよくあります (例: 「program」と「programs」)。この特性を活用するために、完全に接続されたレイヤーを持つ文字埋め込みを介してトークンを表現します。ここで、W(c) は重みであり、文字列は事前定義された最大長 M までパディングされます。完全に接続されたレイヤーの後にレイヤーの正規化も適用します。これらのベクトルは NL エンコーダーにフィードバックされ、サブレイヤーを通じて単語埋め込みと統合されます。

NLエンコーダのニューラルネットワーク構造:

NL エンコーダーは、ブロックのスタック (合計 Nd 個のブロック) で構成されます。各ブロックには、特徴を抽出するための 3 つの異なるサブレイヤー (self_attention、ゲーティング メカニズム、および単語畳み込み) が含まれています。これらについては、次のサブセクションで詳しく説明します。 2 つのサブレイヤー間では、残差接続を使用し、レイヤーの正規化を実行します。

（１）自己注意：自己注意サブレイヤーはTransformerアーキテクチャに従い、マルチヘッドアテンションを使用して長期的な依存情報を取得します。

入力トークンのシーケンス n1、n2、···、nL については、ルックアップ テーブルを介して埋め込み n1、n2、···、nL として表現します。また、位置埋め込みを使用して単語の位置情報をエンコードし、b 番目の Transformer ブロック内の i 番目の単語の位置埋め込みを計算します。ここで、pi,b[·]はベクトルpi,bの次元のインデックスであり、dは次元数です。

Transformer ブロックは、マルチヘッドを通じて非線形特徴を学習し、行列 Y を生成します。マルチヘッド層の計算式は以下の通り（5）で、Hはヘッド数、Whは重みを表す。注意層は各頭部head_tに適用され、(6)によって計算される。ここで、dk=d/H は各固有ベクトルの長さを表します。 Q、K、Vは（7）で計算される。ここで、WQ、WK、WV はモデル パラメーターです。 xi はこの Transformer モジュールへの入力です。最初のブロックの場合、これはルックアップ テーブル埋め込みと位置埋め込みのベクトルの合計、つまり ni+p1,i です。他のブロックの場合、これは下位レベルの Transformer ブロックの出力とこのブロックに対応する位置埋め込みのベクトルの合計です。

（２）ゲーティング機構：自己注意を通じて特徴を計算した後、さらに文字埋め込み情報をマージします。これは、ソフトマックスベースのゲーティング メカニズムによって提供されます。 i番目の単語については、線形変換によってy(self)iから制御ベクトルqiを計算します。文字埋め込みのソフトマックス重みk(c)iは、式2のn(c)iの線形変換によって与えられる。 Transformer出力のソフトマックス重みk(y)iは、y(self)iの別の線形変換によって与えられます。そしてゲートは（８）で計算される。これらは、Transformer レイヤーの特徴 v(y)i と文字埋め込み v(c)i をそれぞれ y(self)i と n(c)i で線形変換するために使用されます。

式5と同様に、このメカニズムの出力はY(gate)=(hi,t)i,tです。ここで、(·)i,tは要素hi,tを持つブロック行列を表します。

（３）単語畳み込み：最後に、２つの畳み込み層がゲーティング機構y(gate)1,...,y(gate)Lに適用され、各トークンy(conv,l)1,...,の近くの局所特徴が抽出される。 y(conv,l)L、ここでlは畳み込み層を表します。 y(conv,l)iは(10)で計算される。ここで、W(conv,l) は畳み込み重み、w=(k-1)/2、k はウィンドウ サイズを表します。特に、y(conv, 0)iはゲーティングメカニズムy(gate)iの出力を表します。これらのレイヤーでは、分離可能な畳み込みが使用されます。その理由は、分離可能な畳み込みの方がパラメータが少なく、トレーニングが容易だからです。最初と最後の単語にはゼロパディングを追加します。これらのレイヤー間では、GELU アクティベーション関数を使用しました。

要約すると、NL エンコーダーには、Transformer の自己注意、ゲーティング メカニズム、および単語畳み込みモジュールの一部が備わっています。自然言語記述は、特徴y(NL)1、y(NL)2、...、y(NL)Lとしてエンコードされます。

AST エンコーダ

生成された部分的な AST の構造をモデル化するために AST エンコーダーを設計しました。私たちのプログラムは文法規則の順序を予測することによって生成されますが、これらの規則だけではプログラムに関する具体的な知識が欠けており、次の規則を予測するには不十分です。したがって、AST エンコーダーは、予測ルールやツリー構造などの異種情報を考慮します。

このようなプログラム固有の情報を組み込むために、まずコードを一連のルールとして表現し、次にアテンション メカニズムを使用してルールをエンコードし、最後にツリー畳み込み層を使用して各ノードとその祖先のエンコードされた表現を結合します。

AST 表現

（１）ルールシーケンス埋め込み：ルール情報をエンコードするためにルールIDを使用します。デコード手順で部分的な AST を生成するために使用されるルールのシーケンス r1、r2、...、rP があるとします。ここで、P はシーケンスの長さを表します。これらのルールは、ルックアップ テーブル埋め込みを介して数値ベクトル r1、r2、...、rP として表現されます。

（２）規則定義エンコーディング：上記のテーブル検索埋め込みでは、文法規則をアトミックトークンとして扱い、規則の内容に関する情報が失われます。この問題を緩和するために、ルール定義のエンコーディングを使用してルールの表現を強化します。文法規則 i: α→β1·βK の場合、α は親ノード、β1·βK は子ノードです。終端文字または非終端文字のいずれかになります。インデックス i はルールの ID です。式2と同様に、ルールの内容を完全に接続された層を介してベクトルr(c)にエンコードし、入力は各シンボルのα、β1、...、βKを埋め込んだルックアップテーブルです。シーケンスも最大長までパディングされることに注意してください。

（３）位置と深度のエンコーディング： ASTデコーダーは自己注意メカニズムを使用するため、文法規則が使用される位置を表現する必要がある。まず、式4の位置埋め込みを採用して、シーケンスr1、...、rPでルールがいつ使用されるかを示します。位置埋め込みはp1(r)、...、pP(r)で表されます。ただし、この位置埋め込みでは、AST 内のルールの位置はキャプチャされません。さらに、このような情報をディープエンベッディングによってエンコードします。シンボルαをルールr:α→β1···βKで展開すると、ルールの深さはその親ノード、つまりαで表されます。このようにして、ルックアップ テーブルの深さ埋め込みの別のシーケンス d1、...、dP を、使用された文法規則のシーケンス r1、...、rP に関連付けます。

AST エンコーダのニューラル ネットワーク構造:

AST エンコーダーもブロックのスタック (合計 N1 ブロック) で構成されます。各ブロックは 4 つのサブレイヤー (自己注意、ゲーティング メカニズム、NL 注意、ツリー畳み込み) に分解されます。ツリー畳み込み層を除き、各サブ層の周囲に残差接続を採用します。各サブレイヤーの後に、レイヤーの正規化を適用します。

（１）自己注意： ASTの情報を取得するために、Transformerのような自己注意層を構築します。ここで、入力はルール埋め込み、位置埋め込み、深度埋め込みの合計、つまりri+di+p(r)iです。自己注意サブレイヤーは、式4、5、6と同じメカニズムを使用して、AST入力の特徴、つまり異なる重みを持つy(ast-self)1、y(ast-self)2、...、y(ast-self)Pを抽出しますが、p(r)iの埋め込み深度を増加させます。

（２）ゲーティングメカニズム：コンテンツエンコーディングルールy(rule)iをTransformerの特徴抽出部分に組み込むことを目指します。式8と9のゲーティングメカニズムを採用します。このサブレイヤーの後、融合された特徴はy(ast-g)1、y(ast-g)2、...、y(ast-g)Pになります。

（３）NL注意：デコードステップでは、入力の自然言語記述を通知する必要があります。これは Multi-headNL によって提供されます。抽出された特徴はy(ast-nl)1、y(ast-nl)2、…、y(ast-nl)Pで表されます。

（4）ツリー畳み込み：上記のサブレイヤーのみを考慮すると、読者がノードの情報とその祖先の情報を組み合わせることは困難になります。通常のシーケンスでは、ノードは祖先から遠く離れている可能性がありますが、構造は緊密です。そのため、従来のTransformerではこのような構造的特徴を抽出することは困難です。ノードの特徴とその祖先の機能を組み合わせます。 AST をグラフとして捉え、隣接行列 M を使用して有向グラフを表します。ノードαiがαjの親ノードである場合、Mji=1になります。すべてのノードが特徴 f1、...、fn によって表されると仮定すると、それらの親ノードの特徴は、隣接行列を乗算することによって取得できます。要約すると、AST デコーダーにはこれら 4 つのサブレイヤーの N1 ブロックがあり、特徴 y(ast)1、y(ast)2、...、y(ast)P を生成します。

デコーダ：

最後のコンポーネントは、生成されたコード情報を自然言語の記述と統合し、次の文法規則を予測するデコーダーです。 AST エンコーダーと同様に、デコーダーでは次のようにブロックのスタック (合計 N2 ブロック) が使用され、各ブロックには複数のサブレイヤーがあります。各サブレイヤーの周囲にも残差接続が採用され、その後レイヤーの正規化が行われます。デコーダーは、展開される非終端記号をクエリとして受け取り、クエリ ノードはルートから展開されるノードまでのパスとして表されます。たとえば、図 1 のノード「Assign」を展開する場合、パスは root、Module、body、Assign になります。このパス内のノードを数値ベクトルとして表します。次に、式 2 のような完全接続層がこれらのベクトルに適用され、パスの出力は qi になります。次に、2 つのアテンション レイヤーを適用して、AST エンコーダーと NL エンコーダーの出力を統合します。最後に、2 つの完全接続レイヤー (最初のレイヤーは GELU アクティベーション関数を使用) が適用され、予測のための特徴が抽出されます。

トレーニングと推論:

デコーダーの最後の層の特徴に基づいて、ソフトマックスを通じてすべての候補単語の中から次の文法規則を予測します。また、自然言語の記述からトークンを直接コピーできるポインター ネットワーク (基本的にはアテンションの一種) も導入します。この場合、生成される文法規則は α→a です。ここで、α は展開される非終端記号であり、a は終端記号です。このポインター メカニズムは、ユーザー定義の識別子 (変数や関数名など) に役立ちます。ソフトマックスまたはポインター ネットワークの具体的な選択は、デコーダーの最後の機能から計算される別のゲーティング メカニズム pg によって行われます。推論は、特殊シンボル snode をルート シンボルに展開する開始ルール start:snode->root から始まります。予測された AST 内のすべてのリーフ ノードが終端記号である場合、再帰予測は終了します。予測時には、サイズ 5 のビーム サーチを使用します。ビーム検索中、無効なルールは除外されます。

要約する

この作業では、TreeGen ジェネレータを使用します。 TreeGen は、Transformer の注意メカニズムを使用して長期的な依存関係の問題を軽減し、AST エンコーダーを導入して文法規則と AST 構造を組み合わせます。評価は、Python データセット Hearthstone と 2 つのセマンティック解析データセット ATIS および GEO で実施されます。実験結果は、私たちのモデルが既存の方法よりも大幅に優れていることを示しています。また、徹底的なアブレーションテストも実施し、モデル内の各コンポーネントが重要な役割を果たしていることを示しています。

謝辞

この記事は、南京大学ソフトウェア学院iSE研究所の2020年度修士課程学生であるCao Zhenfei氏によって翻訳・転載されました。