Seamlessly Integrating Tree-Based Positional Embeddings into Transformer Models for Source Code Representation

生成日: 2026-05-24 02:20

# Seamlessly Integrating Tree-Based Positional Embeddings into Transformer Models for Source Code Representation
- ソースコードを通常の文字列としてだけでなく、抽象構文木（AST）の深さと兄弟順序を使って Transformer に読ませる研究である。
- CodeBERTa-small に小さな追加埋め込みを加え、Masked Language Modeling とコードクローン検出で基準モデルを上回った。
- 既存の Transformer の骨格を大きく変えずに構文情報を入れる点に、実装上の扱いやすさがある。

## 論文の面白いところ

この論文は、コード用 Transformer が見落としやすい「木の中の位置」を、素直な埋め込みとして足し込む。自然言語では単語の前後関係が大きな手掛かりになるが、プログラムでは入れ子、親子関係、同じ階層での並びも意味を持つ。たとえば、同じ `return` でも関数直下にあるのか、条件分岐の内側にあるのかで役割が変わる。著者らは、この差を抽象構文木（Abstract Syntax Tree; AST）の深さと兄弟インデックスで表す。特殊な注意機構を新設するのではなく、CodeBERTa の入力埋め込みに構造の情報を加える構成である。追加されるパラメータは約 78.9 万で、元の CodeBERTa-small の約 0.945% にとどまる。この小ささは、実用上はかなり大事である。構造を使いたいが、モデル全体を作り替えたくない場合に採りやすい方法だからである。結果も派手ではないが一貫しており、論文の主張とよく合っている。

## 問題設定

CodeBERTa などのコード用 Transformer は、トークン列を読んでコードの表現を作る。通常の位置埋め込みは、トークンが何番目に現れるかを表す。これは順序を扱うには十分だが、コードの構造をそのまま表すものではない。プログラムは関数、ブロック、条件分岐、ループなどが階層を作る。抽象構文木は、その階層を木として記述する標準的な表現である。既存の手法には、木構造を明示的に扱う Tree-LSTM やグラフニューラルネットワークもある。しかし、それらは専用の構成を必要とし、標準的な Transformer にそのまま混ぜるのは簡単ではない。この論文の問題設定は、既存の CodeBERTa に近い形を保ったまま、AST 由来の構文位置をどう与えるかである。評価対象は、事前学習に近い Masked Language Modeling と、下流タスクとしてのコードクローン検出である。

## 提案手法

提案手法は Tree-Enhanced CodeBERTa と呼ばれる。まず Tree-Sitter を用いてソースコードから抽象構文木を作る。各トークンには、AST 上で根からどれだけ深い位置にあるかを表す深さと、同じ親を持つ子の中で何番目かを表す兄弟インデックスを割り当てる。これらをそれぞれ学習可能な埋め込みに変換し、通常のトークン埋め込みや位置埋め込みと組み合わせる。組み合わせ方として、単純和、重み付き和、連結して射影する方法の三つを比較している。重み付き和では、単語、通常の位置、トークン型、深さ、兄弟インデックスの寄与を学習中に調整する。論文中の分析では、構造埋め込みは学習の初期に強く効き、その後は単語埋め込みの比重が大きくなる傾向が示されている。また、Tree Attention Mask により、構造的に関係するトークンへ注意を向けやすくしている。骨格となるモデルは 6 層の CodeBERTa-small であり、主要な Transformer 構造は変更しない。

## 結果

実験は CodeSearchNet を用いた Masked Language Modeling と、約 60 万組のコード片ペアからなるクローン検出で行われた。各実験は乱数シード 12345、550、42 の 3 回で平均されている。Masked Language Modeling では、元のモデルの accuracy が 0.8972、F1 が 0.8939 であった。重み付き和の Tree-Enhanced CodeBERTa は accuracy 0.9029、F1 0.8999 となった。損失も 0.44388 から 0.41417 に下がっている。クローン検出では、元のモデルが accuracy 0.9173、F1 0.9172 であった。重み付き和では accuracy 0.9187、F1 0.9186 となり、損失は 0.25836 から 0.21799 に下がった。単純和は Masked Language Modeling では改善したが、クローン検出では元のモデルを少し下回った。連結方式は表現力を増やす一方で、クローン検出では 0.9063 まで落ちた。全体として、構造情報を入れればよいというだけではなく、既存の意味情報とどの比率で混ぜるかが重要であることが分かる。

## 具体例

たとえば、二つの Python 関数があり、どちらもリストの要素を調べて条件に合う値を返すとする。一つ目は `for` ループの中で `if x > 0:` を確認し、その内側で `return x` を実行する。二つ目は見た目のトークン列が似ているが、`return x` が `if` の外に出ており、最後に見た値を返してしまう。この場合、単語や記号の並びだけを見るモデルは、二つのコードをかなり近いものとして扱う可能性がある。Tree-Enhanced CodeBERTa は、`return` が条件分岐の子孫にあるのか、ループ本体の別の子として置かれているのかを AST の深さと兄弟順序から受け取る。モデルはこの構造差を埋め込みに反映し、クローン検出では「同じ機能の実装」とは判断しにくくなる。Masked Language Modeling でも、条件分岐の内側で欠けたトークンを予測するとき、遠くにあるが同じ部分木に属する変数や演算子を参照しやすくなる。間違えやすい点は、字面が似たコードほど意味も同じだと見なしてしまうことである。この論文の手法は、その誤りを完全に消すものではないが、構文上の位置を入力側で明示することにより、モデルが区別に使える手掛かりを増やしている。

Seamlessly Integrating Tree-Based Positional Embeddings into Transformer Models for Source Code Representation

ソースコードを通常の文字列としてだけでなく、抽象構文木（AST）の深さと兄弟順序を使って Transformer に読ませる研究である。
CodeBERTa-small に小さな追加埋め込みを加え、Masked Language Modeling とコードクローン検出で基準モデルを上回った。
既存の Transformer の骨格を大きく変えずに構文情報を入れる点に、実装上の扱いやすさがある。

論文の面白いところ

この論文は、コード用 Transformer が見落としやすい「木の中の位置」を、素直な埋め込みとして足し込む。自然言語では単語の前後関係が大きな手掛かりになるが、プログラムでは入れ子、親子関係、同じ階層での並びも意味を持つ。たとえば、同じ return でも関数直下にあるのか、条件分岐の内側にあるのかで役割が変わる。著者らは、この差を抽象構文木（Abstract Syntax Tree; AST）の深さと兄弟インデックスで表す。特殊な注意機構を新設するのではなく、CodeBERTa の入力埋め込みに構造の情報を加える構成である。追加されるパラメータは約 78.9 万で、元の CodeBERTa-small の約 0.945% にとどまる。この小ささは、実用上はかなり大事である。構造を使いたいが、モデル全体を作り替えたくない場合に採りやすい方法だからである。結果も派手ではないが一貫しており、論文の主張とよく合っている。

問題設定

CodeBERTa などのコード用 Transformer は、トークン列を読んでコードの表現を作る。通常の位置埋め込みは、トークンが何番目に現れるかを表す。これは順序を扱うには十分だが、コードの構造をそのまま表すものではない。プログラムは関数、ブロック、条件分岐、ループなどが階層を作る。抽象構文木は、その階層を木として記述する標準的な表現である。既存の手法には、木構造を明示的に扱う Tree-LSTM やグラフニューラルネットワークもある。しかし、それらは専用の構成を必要とし、標準的な Transformer にそのまま混ぜるのは簡単ではない。この論文の問題設定は、既存の CodeBERTa に近い形を保ったまま、AST 由来の構文位置をどう与えるかである。評価対象は、事前学習に近い Masked Language Modeling と、下流タスクとしてのコードクローン検出である。

提案手法

提案手法は Tree-Enhanced CodeBERTa と呼ばれる。まず Tree-Sitter を用いてソースコードから抽象構文木を作る。各トークンには、AST 上で根からどれだけ深い位置にあるかを表す深さと、同じ親を持つ子の中で何番目かを表す兄弟インデックスを割り当てる。これらをそれぞれ学習可能な埋め込みに変換し、通常のトークン埋め込みや位置埋め込みと組み合わせる。組み合わせ方として、単純和、重み付き和、連結して射影する方法の三つを比較している。重み付き和では、単語、通常の位置、トークン型、深さ、兄弟インデックスの寄与を学習中に調整する。論文中の分析では、構造埋め込みは学習の初期に強く効き、その後は単語埋め込みの比重が大きくなる傾向が示されている。また、Tree Attention Mask により、構造的に関係するトークンへ注意を向けやすくしている。骨格となるモデルは 6 層の CodeBERTa-small であり、主要な Transformer 構造は変更しない。

結果

実験は CodeSearchNet を用いた Masked Language Modeling と、約 60 万組のコード片ペアからなるクローン検出で行われた。各実験は乱数シード 12345、550、42 の 3 回で平均されている。Masked Language Modeling では、元のモデルの accuracy が 0.8972、F1 が 0.8939 であった。重み付き和の Tree-Enhanced CodeBERTa は accuracy 0.9029、F1 0.8999 となった。損失も 0.44388 から 0.41417 に下がっている。クローン検出では、元のモデルが accuracy 0.9173、F1 0.9172 であった。重み付き和では accuracy 0.9187、F1 0.9186 となり、損失は 0.25836 から 0.21799 に下がった。単純和は Masked Language Modeling では改善したが、クローン検出では元のモデルを少し下回った。連結方式は表現力を増やす一方で、クローン検出では 0.9063 まで落ちた。全体として、構造情報を入れればよいというだけではなく、既存の意味情報とどの比率で混ぜるかが重要であることが分かる。

具体例

たとえば、二つの Python 関数があり、どちらもリストの要素を調べて条件に合う値を返すとする。一つ目は for ループの中で if x > 0: を確認し、その内側で return x を実行する。二つ目は見た目のトークン列が似ているが、return x が if の外に出ており、最後に見た値を返してしまう。この場合、単語や記号の並びだけを見るモデルは、二つのコードをかなり近いものとして扱う可能性がある。Tree-Enhanced CodeBERTa は、return が条件分岐の子孫にあるのか、ループ本体の別の子として置かれているのかを AST の深さと兄弟順序から受け取る。モデルはこの構造差を埋め込みに反映し、クローン検出では「同じ機能の実装」とは判断しにくくなる。Masked Language Modeling でも、条件分岐の内側で欠けたトークンを予測するとき、遠くにあるが同じ部分木に属する変数や演算子を参照しやすくなる。間違えやすい点は、字面が似たコードほど意味も同じだと見なしてしまうことである。この論文の手法は、その誤りを完全に消すものではないが、構文上の位置を入力側で明示することにより、モデルが区別に使える手掛かりを増やしている。