PaperWeekly - 代碼表征預訓練語言模型學習指南：原理、分析和代碼－鑽石舞台

©作者 |孫秋實

學校 |華東師範大學

研究方向 |自然語言處理

自從 2020 年 CodeBERT 開了代碼表征預訓練模型（本文稱之為 CodePTM）這個新坑後，在短短兩年的時間內出現了若干個程序設計語言（Programming Language，稱之為 PL，與 Natural Language，也就是 NL 對應）語言模型。它們的共同特點是大部分用於處理 PL（或 Source Code）時所採用的技術是從 NLP 中的 Transformer-based 模型遷移而來，但這些 CodePTMs 又各有特點，從訓練的角度出發即不同的預訓練/微調策略，從數據的角度出發，比如是否使用了代碼的結構信息，即抽象語法樹（Abstract Syntax Tree，本文將其簡稱為 AST）及其變體。

而從架構出發，這些 Code 預訓練模型大致可以被分為以下這三類：

encoder-only：CuBERT、CodeBERT、GraphCodeBERT 等。

decoder-only：CodeGPT、GPT-C 等。

encoder-decoder：PLBART、CodeT5 等。

本文對各個 CodePTM 建模編程語言的思想進行回顧，並簡要分析了一下它們的特色。對文中提到的所有 CodePTMs 的描述主要從背景、預訓練策略、微調策略以及下游任務這幾個角度出發進行分析，考慮到這些模型之間都存在一些共性以及文章篇幅原因，文中略去了一些通用的處理手段和細節，因此對各部分的分析講解詳略不一，不過都保留了建模編程語言最核心的思想。閱讀前需要對 Transformer 有一定的了解。

CuBERT

Learning and Evaluating Contextual Embedding of Source Code.ICML 2020.

https://aclanthology.org/2020.findings-emnlp.139/

CuBERT，即 Code Understanding BERT，和後面提到的 CodeBERT 可被歸為同一個時期的工作，雖是首個提出 Code 預訓練模型的工作，但和 CodeBERT 相比，其影響力較小（在寫這篇文章的時候 CuBERT 引用還沒過百），具體原因個人認為是它僅對 Python 語言進行建模（CodeBERT 同時對 6 種編程語言建模），且它的下游任務和 CodeBERT 相比不太豐富，主要是以代碼理解任務為主。

作者通過 GitHub 採集了 7.4M Python 語言編寫的程序用於 CuBERT 的預訓練，將基於 Transformer 類模型處理自然語言的手段遷移到了編程語言上，使用的模型架構和訓練方式直接照搬了 BERT-Large（24 層 16 個注意力頭），然後使用了一個處理過的ETH Py150 [1]數據集進行微調。與此同時，作者還訓練了一組 Word2Vec embeddings、Bi-LSTM 模型和 Transformer 模型用於比較。

就任務而言，作者構建了一個（在當時全新的）Benchmark，其中包含了：

五個分類任務（具體細節和描述可參考原文附錄）

Variable-Misuse Classification

Wrong Binary Operator

Swapped Operand

Function-Docstring Mismatch

Exception Type

一個定位+修復任務（Variable-Misuse Localization and Repair），也是本文唯一一個非分類任務

就這幾個下游任務而言，可以看到 CuBERT 主要還是在做代碼理解領域的判別式任務，與後續出現的 CodeXGLUE Benchmark 比其在任務的數量和類型上都有局限性，也導致了這個 Benchmark 沒有被廣泛使用。而且，由於它僅採用了 Python 語言，和後面出現的各種 CodePTMs 比局限性也比較大，因此僅做簡單的科普。

CodeBERT

CodeBERT: A Pre-Trained Model for Programming and Natural Languages.Findings of EMNLP 2020.

https://aclanthology.org/2020.findings-emnlp.139.pdf

相比前面提到同期工作的 CuBERT，CodeBERT 的影響力比它大很多。一方面是因為它是多語言（multi-programming-lingual）模型，納入了 6 個編程語言，另一方面是它和 MSRA 自己的 CodeXGLUE Benchmark 配套後在各個下游任務上被廣泛使用。

除了預訓練階段的任務有變化外，CodeBERT 的其他方面與自然語言中的 BERT 模型訓練基本無異（其本質上的一個 RoBERTa），CodeBERT 使用了 bimodal data（即 PL-NL Pairs）進行了預訓練，預訓練數據的來源為CodeSearchNet 數據集 [2]，其中有 Python, Java, JavaScript, PHP, Ruby 和 Go 這六種編程語言的 2.1M bimodal data 和 6.4M unimodal codes（也就是沒有對應 comments 的純代碼），這些數據的來源都是 GitHub 中的開源倉庫，並且後續的很多工作也在預訓練階段用了 CodeSearchNet 數據集。

▲ 表中為CodeSearchNet的數據概覽

CodeBERT 的輸入形式為：

第一段為自然語言文本，第二段為代碼，訓練的數據可分為兩種，即 bimodal data，即 NL-PL Pairs 和 unimodal data，也就是純代碼。

▲ bimodal data

CodeBERT 在預訓練階段的任務有以下兩個：

Masked Language Modeling（MLM），算是 Transformer 類模型的預訓練中最老生常談的任務了，作者將其應用於基於 bimodal data 的訓練。

▲Masked Language Modeling

Replaced Token Detection（RTD），遷移自 ELECTRA，既可以利用 bimodal data 進行訓練，還可以進一步利用 unimodal data（比如是沒有對應自然語言文本的 code），具體細節可以參考 ELECTRA 原文。

▲Replaced Token Detection

實驗部分做了 Natural Language Code Search，個人認為文中沒有添加更多下游任務是受到 EMNLP 的篇幅限制，使用 CodeBERT 可以在多個下游任務，如 Clone detection（克隆檢測）、Defect detection（缺陷檢測）、Code summarization 等上得到出色的結果，具體可參考CodeXGLUE [3]，如下圖所示：

▲ CodeXGLUE: A Machine Learning Benchmark Dataset for CodeUnderstanding and Generation

從 CodeBERT 開始，後續的 CodePTMs 就全部繼承了對多個 PL 的支持，不過 CodeBERT 完全使用了建模自然語言的手段來為 Code（或是說 NL-PL Pairs）做預訓練，忽視了代碼的一個很大的特性，那就是結構信息，具體而言就是在編譯器進行語法分析階段生成的抽象語法樹（Abstract Syntax Tree，本文稱之為AST），緊跟着 CodeBERT 的 GraphCodeBERT 立刻填上了這個坑。

GraphCodeBERT

GraphCodeBERT: Pre-training Code Representations with Data Flow.ICLR 2021.

https://openreview.net/forum?id=jLoC4ez43PZ

看名字就知道這是 CodeBERT 的後續工作，主要想法就是為 CodeBERT 添加代碼的語法信息，使 CodePTM 可以顯式學習代碼的結構信息。

GraphCodeBERT 基於數據流學習代碼的表徵，如下圖所示

▲ GraphCodeBERT所使用的數據流

數據流的獲得分為以下幾個步驟：

通過語法分析工具獲得 AST，原文中使用的工具是 tree-sitter。

從 AST 中提出變量，構成一個由變量組成的序列。

從 AST 中抽取變量之間的依賴關係，文中稱之為「value comes from」，構造數據流圖。

▲GraphCodeBERT的預訓練

GraphCodeBERT 在模型預訓練階段額外提出了兩個在當時較為新穎的訓練任務

Edge Prediction，即數據流圖邊預測，通過預測數據流圖的邊學習代碼的結構信息。

Node Alignment，即變量對齊，具體而言是學習數據流圖中的某個 node 來自輸入代碼中的哪個 code token。

將它們和從 CodeBERT（或是 BERT or RoBERTa）繼承下來的 MLM 任務一起優化。考慮到 AST 是一種圖結構，為了讓 Transformer 能適應其與一般序列結構的差異，作者修改了其注意力機制，主要是通過調整 Attention Mask 縮小感受野。

作者將其稱為 Graph-Guided Masked Attention，其中 E 代表的是數據流圖的邊，E' 代表的是數據流圖的節點和代碼的對應關係邊。

就下游任務而言，GraphCodeBERT 文中主要完成了 Natural Language Code Search、Clone Detection、Code Translation 和 Code Refinement 這幾個任務，它同樣適用 CodeXGLUE Benchmark 中的其他任務，比如 Code Summarization 等。

GraphCodeBERT 相較於前作 CodeBERT 解決了 CodePTM 只學習自然語義，而不學代碼結構/語法的問題。但細心的讀者或許能發現，學習數據流相較於學習 AST 本身有相當的信息損失，這也為之後的 UniXcoder 挖了一個小坑。

GPT-C

IntelliCode Compose: Code Generation Using Transformer.FSE/ESEC 2020.

https://arxiv.org/abs/2005.08025

GPT-C 是為了代碼補全（Code Completion）這個任務而設計的，作者認為之前的的代碼補全工作有兩點不同。

根據上文的 token 來預測下個 token，沒有將代碼的全文環境納入考慮；

多語言效果不佳。

▲ Code Completion

作者提出的 GPT-C 是 GPT-2 模型的變體，在一個大規模、無監督、多語言的數據集上從零開始訓練。基於 GPT-C，作者構建了一個代碼補全 Framework，稱之為 IntelliCode Compose，並對多種編程語言進行建模。作者將 Sequence decoding 的過程視為對樹的搜索，搜到出現目標 token 為止。

▲Sequence decoding

雖說是多語言，但是使用的是 Python, C#, JavaScript 和 TypeScript，和 CodeXGLUE 不同且少了兩個語言。就多語言模型的訓練而言，作者提出了四個訓練的策略

Language-agnostic baseline，即忽略掉編程語言的不同構建一個 baseline 多語言模型。

Language-type embedding，即加入一個向量來表示每種編程語言，和 token embedding 等相加。

Language-specific control codes，每個輸入的訓練樣本前拼接一個"lang ∗ remaining token sequence」字符串，∗即為編程語言。

add a programming language classification task during model pretraining，即在預訓練階段加入一個分類編程語言的任務。

隨後在下游任務上驗證了這幾個方案的效果，證明了 Language-type embedding 和 control codes 方案在對各編程語言建模方面的有效性。

▲multilingual modeling approaches based on GPT-C

在文末，作者考慮到了模型的推理開銷問題，還上了一個知識蒸餾，並且還討論了一下模型基於 K8S 和 VS Code 的部署應用等問題。

Code-GPT

Code-GPT 是在 CodeXGLUE 中被提出的，沒有單獨成文，不要和 GPT-C 搞混了。作者實現它的目的是為了 code completion 和 text-to-code generation 任務。它就是一個由 Code 訓練，與 GPT-2 完全同架構的 12 層 Transformer Decoder 模型，不過 MSRA 的研究者實現了兩個版本。

Pretrained from scratch：隨機初始化，從零訓練；

CodeGPT-adapted：先使用一個 GPT-2 作為起始點，再持續使用 Code 進行訓練，作者將這個方法稱為「domain-adaptive」。

更詳細的內容可以參考 CodeXGLUE 原文的 4.2 節，作者在 Huggingface 提供了CodeGPT-small-java [4] 和CodeGPT-small-java-adapted [5] 這兩個 checkpoints，正常地使用 transformers 庫加載就能使用了。

PLBART

Unified Pre-training for Program Understanding and Generation.NAACL 2021.

https://arxiv.org/abs/2103.06333

顧名思義，就是應用於編程語言的 BART，參考了文章：BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension（ACL 2020）[6]。

先簡單說下 BART，它吸收了 BERT 模型的雙向編碼器和 GPT 中的單向 left-to-right 解碼器這二者的優點，因此比 BERT 更適合文本生成的場景（BERT 因 Pre-training 階段和生成式下游任務差異比較大，因此被認為不適合做 NLG 相關任務），而它相比 GPT，也多了雙向上下文語境信息（GPT 是單向建模）。除此之外，相比 BERT 的 Token Masking，BART 對 Encoder 端採用了更加複雜的Noise。

基於對 BART 的知識，理解 PLBART 並不算困難，其使用了和 BART-base 相同的架構，唯一結構上的不同的是它在 Encoder 和 Decoder 的頂部添加了一個額外的 LayerNorm。在 Noise 策略方面，PLBART 使用了 token masking, token deletion 和 token infilling 這三種策略，與 BART 相比少了 Sentence Permutation 和 Document Rotation 這兩個任務，這些任務的細節都可以參考 BART 原文。

在微調下游任務方面，作者將 PLBART 的下游任務分為兩塊。

首先是 Sequence Generation，又可細分為三個任務，可參考下圖：

▲ PLBART: Sequence Generation

其次是 Sequence Classification：將序列尾部的 special token 餵給線性分類器用於預測，與 BERT 等模型的分類區別不大。

實驗與比較方面，作者先指定了 baseline 模型，並將其分成了兩種：

Training from Scratch，作者用下游任務的數據集從零開始訓練了 LSTM + Attention 以及一個 Transformer。

Pre-trained Models，作者挑選了 RoBERTa、CodeBERT、GraphCodeBERT、GPT-2、CodeGPT（-adapted）。

具體的實驗部分做了 Code Summarization、Code Generation、Code Translation 這三個生成式任務，效果自然是好的，在 Classification 方面做了兩個任務：clone detection 和 vulnerability detection，在後者上 PLBART 不如基於 AST 的模型。

CodeT5

CodeT5: Identifier-aware Unified Pre-trained Encoder-Decoder Models for Code Understanding and Generation.EMNLP 2021.

https://aclanthology.org/2021.emnlp-main.685/

文中對 CodeT5 的描述是：a unified pre-trained encoder-decoder Transformer model that better leverages the code semantics conveyed from the developer-assigned identifiers，即一個能更好地利用代碼語法信息（形式是 identifier，即標識符）的統一預訓練 Transformer 模型。在開始之前，和 PLBART 一樣，先簡單說下 Google T5 模型。T5 的名字來源是 Text-To-Text Transfer Transformer，顧名思義 T5 把所有的 NLP 問題統一歸納為了 Text2Text 任務，用來做 NMT、QA、文本摘要和文本分類任務。

▲ Text-To-Text Transfer Transformer

對比下 T5 原文，可以發現二者的核心思想還是非常類似的，作者將 CodeT5 歸納為 a pre-trained encoder-decoder model that considers the token type information in code，細心的玩家可能發現了，前面提到的 CodeBERT 為首的 BERT 類模型和 CodeGPT 為首的 GPT 類模型，僅含有 Encoder 或 Decoder，而非完整利用一個 Transformer 架構來處理代碼。因此 CodeT5 最大的賣點即第一個 unified encoder-decoder CodePTM，可以理解為完全使用了 Transformer 的兩個部分。

▲ CodeT5 Illustration

此外，除了使用 T5 的架構外，作者使用了以下兩個方案來更好地利用代碼結構特性：

使用了代碼的標識符信息，提出了 Identifier-aware Pre-training，是一個與 T5中Masked Span 類似的目標，簡而言之就是隨機 mask 掉任意長度的 Span，然後讓 decoder 去預測。

利用了代碼地 Comments（自然語言注釋）信息，作者稱之為 Bimodel Dual Generation，讓自然語言和源代碼的表徵可以對齊，幫助緩和 Pre-training 和 Fine-Tuning 階段的差距。

▲ Identifier-aware Pre-training

文章發表時在 CodeXGLUE Benchmark 的若干任務上取得了 SOTA 效果。

UniXcoder

UniXcoder: Unified Cross-Modal Pre-training for Code Representation.ACL 2022.

https://arxiv.org/pdf/2203.03850.pdf

這是當今（2022.7）的 SOTA 模型，參考了 NIPS 2019:Unified Language Model Pre-training for Natural Language Understanding and Generation [7]。

文章選擇了五個任務，分為以下三類：

代碼理解任務：Clone Detection 和 Code Search

代碼生成任務：Code Summary 和 Code Generation

自回歸任務：Code Completion

本文很重要的一個賣點就是更全面地利用了 AST 提供的代碼結構信息。文章開頭講過，AST 一般會被表示為一個 Tree 結構，不能直接作為 Transformer 類模型的輸入，回憶一下前面提到的 GraphCodeBERT，作者在以損失相當一部分信息的情況下讓模型學習 AST 的數據流。為了能更加有效地利用 AST，因此 UniXcoder 地作者構建了一個 one-to-one 的 mapping function，將 AST 轉為一個序列結構（flattened token sequence），然後和Code Comments一同編碼，對這個 mapping function 的有效性的證明在文章的附錄中。

模型結構方面，UniXcoder 的一個賣點就是一個統一的，可以同時兼容 Encoder-Only，Decoder-Only 和 Encoder-Decoder 三種模式的 CodePTM，相當於給模型添加了「開關」，來決定採用什麼模式處理任務，用白話講，就是通過使用三種不同類型的自注意力 Mask 策略來控制模型的行為。

▲ UniXcoder的架構

既然同時能擁有三種模式，那麼自然會有更多預訓練任務，如下所示：

Masked Language Modeling（MLM），算是基本操作了。

Unidirectional Language Modeling（ULM），用於訓練 decoder-only 模式，幫助完成自回歸任務，對應的是右上三角 masking。

Denoising Objective DeNoiSing（DNS），可參考 BART 和 T5，用於訓練 encoder-decoder 模式，幫助完成生成任務，參考架構圖中的 encoder-decoder 部分。

除了上面這些任務以外，作者還提出了 Code Fragment Representation Learning。

▲ Code Fragment Representation Learning

其中包含了 Multi-modal Contrastive Learning（MCL）和 Cross-Modal Generation（CMG）這兩個任務。前者採用了一個對比學習損失，後者是使用了一個統一的自然語言描述（comment），文中使用了 fulcrum，即支點這個詞，讓模型學習到的代碼表征在不同語言之間的對齊。

還需注意的一點就是，UniXcoder 在預訓練和微調這兩個階段中的輸入形式有所不同，由於引入了 Flattened AST，AST 展開後的序列中被引入了大量額外的 tokens（70% longer）會導致額外的開銷。因此，在微調階段 UniXcoder 僅使用 AST 的葉子節點，為了緩解這個 gap，在預訓練階段作者設置了 0.5 的概率隨機丟棄輸入序列中的非葉子節點。

除了 Clone Detection、Code Search、Code Summarization 和 Code Completion 等任務上表現較好外，UniXcoder 還提供了一個新任務：zero-shot code-to-code search，即在 zero-shot 的情境下，通過一個源代碼的 query 在 candidates 集合中尋找語義相同的源代碼，該任務使用的數據集是 CodeNet [8]，用來衡量訓練所得的 code fragment embeddings 的效果。

▲ zero-shot code-to-code search

相關代碼整理

CuBERT：

https://github.com/google-research/google-research/tree/master/cubert

CodeBERT、GraphCodeBERT 和 UniXCoder

MSRA 提供了 CodeBERT、GraphCodeBERT 和 UniXCoder 在下游任務微調時可用的代碼，在倉庫：

https://github.com/microsoft/CodeBERT

但沒有提供預訓練階段的實現（CodeBERT 和 UniXCoder 在預訓練階段都使用了 16 張 32GB NVIDIA Tesla V100 實現），使用時通過 transformers 加載 checkpoints 就可使用。

此外，huggingface 還提供了一個經濟適用版的 CodeBERT 模型：

https://huggingface.co/huggingface/CodeBERTa-small-v1

CodeGPT

與上述三個 MSRA 提供的模型一樣，CodeGPT 仍然是提供了可通過 transformers 加載 checkpoints，即：

https://huggingface.co/microsoft/CodeGPT-small-java-adaptedGPT2

CodeT5：

https://github.com/salesforce/codet5

PLBART：

https://github.com/wasiahmad/PLBART

總結

上述對 CodePTMs 相關的內容大致上是六月中下旬趕一個文章時調研文獻後總結的筆記，然後補充了點鏈接和細節，實驗部分寫的比較簡略，是因為如果每個模型的實驗都要全講的話篇幅就太長了，而且這些任務都大差不差，每個模型都講一遍冗餘會比較多，之後可能會在其他文章補充。

此外，有一些影響力比較小或者 Task-specific 的工作可能沒完全覆蓋到。總而言之，不論是 CodeBERT 開的新坑還是今天的 SOTA 模型 UniXcoder，MSRA 在這個領域還是完全 dominant 的存在。對於 CodeBERT 和 GraphCodeBERT 為首的大模型，復現預訓練階段的成本很高，不適合平民玩家，而且今年五月的 IJCAI 22 Survey Track 連 CodePTMs 的 Survey 工作都已經出了（Deep Learning Meets Software Engineering: A Survey on Pre-Trained Models of Source Code [9]），可能短時間內出革命性的新模型的可能性不大，而且 Code 領域使用的這些方法終究還是跟着 NLP 走的，需要 NLP 提出新技術後 Code 領域才有跟進的可能。

個人感覺接下來在這個相對較小但很卷的領域的研究熱點可能會慢慢向可解釋性和模型分析（Analysis of Models & Interpretability）方面轉移，最近還研讀了一些 22 年新出的 Probing CodePTMs 的文章，之後再補充。