Zenn
📜

[論文輪読] A Deobfuscation Pre-Training Objective for ...

2021/03/23に公開
機械学習
transcompiler
tech

A Deobfuscation Pre-Training Objective for Programming Languages

  • 論文
  • Baptiste Roziere, Marie-Anne Lachaux, Marc Szafraniec, Guillaume Lample
  • Paris Dauphine University, Facebook AI Research
  • 2021/2/16 publish

概要

ML初心者なのであちこちおかしいところがきっとあります。

  • ソースコードの難読化解除を教師なし学習で学習させたモデルで実現した
  • MLMをプログラミング言語に特化させたモデル

キーワード

  • MLM: Masked Language Modeling

1. Introduction

コード難読化は人間が理解しにくくするためであったり、機能を保ったままサイズを小さくするために行われる。現代においては主にコードの隠蔽化やマルウェア検出から逃れるためであったり、JavaScriptなどのコードを圧縮して通信量を削減するために行われる。例えばCのコードの場合変数名の変更などの不可逆な変換も行われる。(コンパイル&デコンパイルのことを指していると思われる)
難読化されたファイルの全体を解読するには、人が行う場合はコードの全体から何をしているかを把握して、細部の難読化を解いていくことになる。この作業には時間がかかるタスクである。さらに、適切な変数名・関数名を復元するのはプログラムが何をしているのかを解析するよりも難易度の高いタスクとなる。

この論文ではこのようなタスクを行うために、seq2seqモデル(sequence to sequence model)をベースとした難読化された関数を元の形に戻すモデルを提案している。

事前学習にはMLM(Masked Language Model：文章の穴埋め予測をするモデル)を用いる。実際にはランダムにマスクがされるため、複数回出現する変数の片方だけがマスクされるといったようなことが起きる。（これは本質とは異なり）前後のコードから容易に予測できることになるため、この論文では同じ変数名・関数名をF0, F1, V1, V2といったスペシャルトークンで置き換える。

この論文のcontributionは

  • DOBFという複数言語に対応した難読化解除器を提案
  • DOBFがcode search, code summarization, unsupervised code translationといったタスクで単純なMLMよりも優れていることを示す
  • DOBFのモデルが完全に難読化されたソースファイルを難読化解除できることを示す

2. Related work

割愛

3. Model

3.1. MLM for Programming Languages

無数の文献でプリトレーングの手法が紹介されている（個々のトークンをマスクすべきかどうか、どの部分をマスクするかなど）。これは自然言語処理で活発に検証が進められている。プログラミング言語は自然言語よりもはるかに構造化されているので、トークンの予測は自然言語のそれよりもはるかに容易なタスクである。

例

図1（上図）ではPythonのキーワードや構文を置き換えている。[, while, =, if, ), returnなどのシンボルは回復がしやすく、完全な正確さで予測することができ、モデルもすぐに学習する。そのため、MLMにとってプログラミング言語の文法構造はとても簡単なものである。

しかし、プログラムの難読化解除というタスクにおいてはそれだけでは不十分で、構造だけでなくプログラムの意味を理解することが必要になる。本論文で提案されるDOBFはこのタスク達成のために、MLMよりもより深くプログラムの意味を理解する。

3.2. Deobfuscation Objective

DOBFはMLMよりもプログラミング言語の構造に特化させたもの。クラス名・関数名・変数名はスペシャルトークンで置き換えることでコードスニペットを難読化して元の名前を復元するモデルをトレーニングする。識別子が選択されると、コード内の全ての同じ変数名・関数名などが同じトークンで置き換えられる。MLMと異なりプログラムの予約語は置き換えない。例えば図1の場合、MLMではマスクされている[と置き換えられていない[がある。反対に、queueはMLMでは一部しかマスクされていないのに対し、DOBFでは全てV3でマスクされている。

識別子は確率P_{obf}\in[0,1]で置換される（ただし、かならず1つ以上置換されるようにする）。つまり、P_{obf}=0のときには識別子1つが置換され、P_{obf}=1のときにはファイル内全ての識別子が置換される。DOBFではコードの予約語や演算子に対しては置換処理を行わないので、置換を行ってもコードの実行結果は同じになる。（図1は幅優先探索のコードをP_{obf}=1で置換している）

良い結果を得るためにはモデルがコードの意味を深く学習している必要があり、事前学習をすることの意味である。MLMのインプットではランダムにマスクをされているので、マスクのされ方によってはコードの本質を理解していなくてもマスク元が推測できる（例えば図1のMLMの４行目ではqueueが呼ばれているが、これを元に3行目のマスクがqueueであることが容易に推測される）。

3.3. Implementation

難読化解除タスクにおいては教師あり（なし？）機械翻訳のようにseq2seqモデルが学習される。

このモデルは推論時に難読化された識別子を持つコードに対して意味のあるクラス名・関数名・変数名を提案することができる。難読化されたクラス・関数・変数はCLASS_0...CLASS_N, FUNC_O...FUNC_N, VAR_0...VAR_Nのように置き換えられる。出力である辞書ではデリミタ記号|で各エントリが区切られる。

4. Experiments

4.1. Deobfuscation

P_{obf}=0とP_{obf}=1の2パターンで難読化解除タスクを行いモデルを評価する。

P_{obf}=0の場合

この場合では全体で一つの識別子のみ難読化されている。モデルは残りのコードをコンテクストとして難読化された識別子に対する意味のある名前を提案しなければならない。（PyCharmやIntelliJがシンプルなルールを使ってこのタスクを行える）

P_{obf}=1の場合

一般的に使われる難読化器はたいてい複数の変形を行う。そのうちの一つとして、全ての識別子の名前を短くて意味のない名前（例えばa, b, c）に置き換える。P_{obf}=1の置換はこの変形に相当する。モデルの性能を測るために、この変形を施したコードをモデルで元に戻させる。

評価基準

  • 元の識別子との完全マッチ割合
  • サブトークンのスコア
    • 正解率：予測した値が実際にそうである割合
    • 再現率：実際に正であるもののうち、正であると予測されたものの割合
    • F値：再現率と適合率の調和平均

モデルの評価には難読化前の識別子を用いる。元の識別子との完全マッチを真値として復元に成功した割合（パーセンテージ）で評価を行う。
また、先行研究にならって、サブトークンスコアも用いる。（サブトークンスコアとはより柔軟な評価軸で、大文字小文字を区別しないサブトークンを検索するための精度・リコール・F1スコア。。。。とあるが、よくわかってない。）

オリジナルの識別子大文字小文字区別なしのサブトークンに対してそれぞれを計算する。それぞれのトークンはアッパーケースのキャメルケースとアンダースコアのスネークケースに分解される。例えば、decoderAttentionの場合にはdecoder_attentionとattentionDecoderにも完全一致する（順番変わるのはありなのか・・・・？）。attentionであれば正解率は高いが再現率は0.5なので、F値は66.7となる。crossAttentionDecoderであれば正解率は2/3で、F値は80.0となる。

サブトークンの正解率・再現率・F値の平均は全ての復元されたトークンに対して計算する。

（ここの取扱がちゃんと理解できてない・・・）

4.2. 各タスク向けFine-tuning

プリトレーニングモデルとしてのDOBFを評価するため、DOBFをTransCoderとCodeXGLUEの3つのタスク向けにファインチューニングする。ここではJavaとPythonのみを対象として行う。

CodeXGLUE Clone Detection

2つのコードスニペットが意味的に等価であるかどうかをモデルが予測する二値分類問題。評価にはF値を用いる。このモデルは、シングルエンコーダとクラシフィケーションレイヤで構成されている。入力は2つのコードスニペットで、モデルに与える前にそれらを連結する。このタスクはJavaで行う。

CodeXGLUE Code Summarization

モデルはコードスニペットが与えられるとそれに対応するドキュメントを自然言語で生成するように学習されている。アーキテクチャはBLEUスコアで評価されたseq2seqモデル。（よくわからん）データセットにはJavaとPythonが用いられている。

自然言語によるコード検索クエリが与えられると、モデルはコードスニペットのコレクションの中から最も意味的に合致するコードを検索する。これはMean Reciprocal Rank(MRR)という指標を用いて評価されるランキング問題。モデルは2つのエンコーダで構成されている。自然言語のクエリとコードは別々にエンコードされて、エンコーダの最終段の第一隠れ層の状態のドット積を計算する（よくわからん）。このタスクはPythonで利用できる。

TransCoder

前回の輪読で紹介した論文で評価に用いられているコードの翻訳タスク。ビームサイズ1と10でおこなう。このタスクは用意されたユニットテストが通過したかどうかで評価されいた（確か・・・）

4.3. 実験の詳細

Model Architecture

アテンションメカニズム有りのseq2seqモデルを使用する。このモデルはトランスフォーマーアーキテクチャを用いたエンコーダとデコーダで構成されている（このあたりの関係がよくわかってない。この論文以前の問題かな）
基準（CodeBERT, TransCoder）との公平な比較を行うために異なる2つのモデルをトレーニングする。

  • 12層、12のアテンションヘッド、768の隠れ次元
  • 6層、8のアテンションヘッド、1024の隠れ次元

Traning dataset

  • GoogleのBigQueryでGithubからpublicなコードを持ってくる
  • 対象はPython・Javaで書かれたコード
  • 重複しているものは削除
  • 各ファイルを難読化して、識別子を得る
    • Python：19GB
    • Java：26GB

Traning details

DOBFは難読化されたファイルを識別子のリストに変換するように学習する。

  • GPUごとに3000のトークンからなるJava・Pythonのバッチを交互に使用する
  • Adam optimizerと逆平方根学習率スケジューラを用いてDOBFを最適化する（詳細は調べてない）
  • モデルはPyTorchで実装
  • 32台のV100GPUで学習
  • 学習を高速化するためにfloat16命令を仕様
  • 初期化は異なる2つの方法で行う
    • 0から学習
    • Python-Java MLMによる学習
  • DOBFを3つの異なる難読化確率（0, 0.5, 1）で学習させる
    • 検証データセットで計算されたサブトークンのF1スコアの平均値を用いて最適なものを選択する。

Fine-tuning details

エンコーダとデコーダを備えたseq2seqモデル、2つのエンコーダを備えたモデル、1つのエンコーダを備えたモデルなど、目的に応じて様々なモデルアーキテクチャを検討する。比較を公平に行うため、元の論文と同じアーキテクチャ、GPU数、バッチサイズ、オプティマイザでモデルを学習。

CodeXGLUEでは、ファインチューニングの際に使用した学習率パラメータにタスクが非常に敏感に反応するめ、5x10^-6〜10^-4までの5つの学習率パラメータでグリッド検索を行い、検証データセットで最適なパラメータを選択。

TransCoderでは元論文と同様の10^-4の学習率を使用している。

実験結果

5.1. Deobfuscation

P_{obf}=0の評価を行う場合でも、P_{obf}=0で学習すると、各入力シーケンスに対して単一の変数のみ生成するようにしかモデルが学習しないので、P_{obf}=0.5よりも（学習？）効率が悪くなってしまう。ただし、P_{obf}=0.5での学習はモデルにコードの意味も理解させるより難易度の高いタスクなので、学習が難しい。したがってP_{obf}=0で学習する場合でもコードの構造学習という点においては優れているので有効。

表2

実験結果（論文の図2）から、DOBFが最もよく復元できたのはP_{obf}=1のときで、45.6%であった。また、全ての場合においてDOBF with MLMのプリトレーニング（MLMのモデルを元にDOBFでトレーニング）はパフォーマンスを向上させた。

図2

図2（上図）はDOBFで関数を完全にリカバリーした時の図。DOBFは関数の目的を理解して入力元に近しい変数名を予測している。

図3

図3（上図）はDOBFがPythonで書かれた行列操作する関数の実装に対して変数名を提案している例。（論文のAppendixにある図5, 6, 7にPythonの変数名、Javaの変数名・関数名の例もある）

図8

Appendixの図8では深さ優先探索・幅優先探索のコードを並べて比較し、微妙な違いも理解した上で変数の提案をしているらしい。（ちゃんと２つのアルゴリズム勉強しなきゃ・・・・）

5.2. Downstream tasks

ファインチューニングするにあたり、プリトレーニングではP_{obf}=0.5とP_{obf}=1では非常に似た結果を得たため、以降P_{obf}=0.5のみ取り上げる。ベースラインとしてランダムに初期化されたモデルとMLMのみでプリトレーニングしたモデルを用いる。また、CodeXGLUEタスクについてはCodeBERTもベースラインとして用いる。

表3

上表（論文の表3）はDOBFのみでトレーニングしたモデルとMLMを初期値としてDOBFでトレーニングしたモデルの比較。ランダムに初期化したモデルは与えられたタスクに対してどれだけプリトレーニングが影響を与えるかを測る指標となる。プリトレーニングは特にNLCSタスク（コード要約タスク）において重要であった。MMRが0.025から0.383までMLMのプリトレーニングにより向上している。（突然指標が出てきた・・・なんぞ）

MLMのベースラインとCodeBERTの最も大きな違いは、CodeBERTは

  • CodeBERTは異なるデータセットでトレーンングされていること
  • 追加のRTDオブジェクトを使用していること
  • RoBERTaモデルで初期化されていること

コード要約とNLコード検索は自然言語を含むため、コメントを含むCodeBERTのデータセットの恩恵をウケる可能性があるが、このタスクではより単純なデータセットを使用して非常に近い結果が得られた。しかし、Clone Detectionにおいてはそのパフォーマンスに及ばなかった。また、RoBERTaでMLMモデルを初期化することも試みたが、downstreamタスクのパフォーマンスに大きな影響を与えることはなかった。

スクラッチでトレーニングしたDOBFとMLM+DOBFはいずれもdownstreamタスク全てに対して到達水準（CodeBERTとMLM）を満たす能力を獲得した。難読化解除では事前学習タスクとしてすでに効果的で、ほとんどのタスクでMLMと同等の結果をもたらし、クローンの検出ではより効果的である。MLM+DOBFのモデルは全てのdownstreamタスクに対してMLMを凌ぐ性能を獲得している。特にNLコード検索で大きな改善が見られた。

TransCoderに対しては、MLM+DOBFによってビームサイズ10のPythonからJavaへの変換で2.7%、JavaからPythonへの変換で5.9%、MLMモデルの計算精度を向上させた。

図4

図4（上図）では、テストセットの計算精度はEpoc数にかかわらずMLM+DOBFのほうが高いことが示されている。また、コード検索やコード要約においてもMLM+DOBFはCodeBERTに大差をつけており、これらのタスクで効果的なモデルを学習するために自然言語に沿ったプログラミング言語データ（多分コメント有りのコードの意味）は必要ないことがわかる。

MLM+DOBFではほとんどのタスクでDOBFとMLMの両方よりも高いスコアが得られていて、MLMとDOBFが補完関係にあることがわかる。

6. 結論

新しい難読化解除器(DOBF)を提案した。DOBFは難読化されたコードの回復・関連する識別子の提案・プログラミング元g関連タスクの変換モデルの事前学習という3つの目的に使用できることが示された。DOBFは自然言語に頼ることなく能力を獲得できていて、クローン検出・コード要約・自然言語コード検索・教師なしコード翻訳といったタスクにおいて、CodeBERT・MLMの事前学習よりも優れている。

(ここからは翻訳したけどよくわからん)
これらの結果は、DOBFがソースコードの特殊な構造を利用して、特に効果的な方法で入力シーケンスにノイズを加えることを示しています。ソースコードに適合した他のノイズ関数やサロゲート目標があれば、さらに性能が向上する可能性があります。例えば、与えられた変数の型、メソッドのシグネチャ、破損したコードの修復などをモデルに学習させることができます。ソースコードで事前学習したモデルは、構造化されたノイズの恩恵を受けているので、この知見を自然言語にも適用できるかどうかは興味深いところである。自然言語は曖昧ではあるが、その根底には構造が存在している。プログラミング言語の抽象的な構文木ではなく、文の構成要素や依存関係の構文木を活用することで、自然言語の前処理目的をよりよく設計できるかもしれません。

Discussion