Towards Compositional Generalization in Code Search (2022)

Abstract: We study compositional generalization, which aims to generalize on unseen combinations of seen structural elements, for code search. Unlike existing approaches of partially pursuing this goal, we study how to extract structural elements, which we name a template that directly targets compositional generalization. Thus we propose CTBERT, or Code Template BERT, representing codes using automatically extracted templates as building blocks. We empirically validate CTBERT on two public code search benchmarks, AdvTest and CSN. Further, we show that templates are complementary to data flow graphs in GraphCodeBERT, by enhancing structural context around variables.

## introduction we study compositional generalization in the context of code search, which aims to generalize on unseen combinations of previously seen elements, or building blocks @xcite . following the categorization in the nlp domain @xcite , figure 1 motivates such a goal in the code search setting, where a model is expected to generalize on an unseen code, which is a combination of a seen structure with other lexicons (lexical generalization). moreover, the model should also generalize on another case when an unseen code is a combination of seen structures, though red and blue colored structures were only independently observed while its combination was not (structural generalization). existing code search models such as code-bert @xcite regarding code as an unstructured sequence (figure 2a ) often fail to generalize when the lexical perturbation without changing structures and labels is applied in test codes, known as advtest evaluation @xcite . towards better generalization, syncobert (wang for _ in _: if _: print(_) … for key in dict.keys(): if dict[key] is not none: print(dict[key]) … unseen … for i in range(10): if @xmath0: print(i) … seen (a) lexical generalization. ... try: a /= b except zerodivisionerror: @xmath1... try: _ except _: raise _ while _: _ if _: break … try: while true: @xmath2: break except typeerror: raise bar() … … while @xmath3: @xmath4: break … (b) structural generalization. et al., 2021b) extracts an abstract syntax tree (ast; figure 2b ) to incorporate syntactic structure into the model. on the other hand, graphcode-bert @xcite extracts a data flow graph (dfg; figure 2c ), to represent variable relations within code. from structure-aware pretraining objectives, these two models significantly enhance their accuracy in code search, both in original and adversarial settings. meanwhile, utilizing composable building blocks, like the colored structures in figure 1 for code search was not studied. inspired by the compositional generalization work in nlp domain, we study how to represent composable building blocks targeted to improve compositional generalization. specifically, we consider subtrees of ast as such building blocks, which we call templates (figure 2d ). as illustrated in figure 1 , templates share across codes, as a code is a combination of templates. however, it requires an additional non-trivial challenge of mining such a meaningful block. in this paper, we suggest how to mine templates and propose ctbert, code template bert, pretrained with templates for compositional generalization. specifically, we group frequently adjacent ast nodes into subtrees, filter uncommon or meaningless subtrees, and regard the remaining ones as templates. then, we pretrain ctbert with templates, using our proposed objective teaching models to recognize structural building blocks. to confirm the effectiveness of ctbert towards compositional generalization in code search, we validate ctbert by fine-tuning on a largescale dataset csn @xcite then testing on an adversarial test set ad-vtest @xcite . advtest is more appropriate for evaluating compositional generalization than csn as 1) variable and function names in test codes are normalized (testing lexical generalization) and 2) test time codes are unseen combinations of seen structures (testing structural generalization). our experimental results show that ctbert enhances its compositional generalization with templates, as shown by the performance gains on ad-vtest, while performing comparably on csn. further, we show that ctbert can straightforwardly be combined with dfg objectives from graph-codebert, achieving orthogonal and complementary performance gains. all our implementations and data including the mined templates will be publicly released 1 . ## ctbert our goal is to improve the model's compositional generalization capability, by pre-training with automatically mined templates. ## experiments to confirm whether learning with code templates contributes to compositional generalization in code ## conclusion this paper studied compositional generalization in code search, and proposed ctbert. our contributions are three-fold: 1) we proposed templates that are composable building blocks for code, and pretrained ctbert with templates; 2) we showed ctbert improves compositional generalization; 3) we empirically confirmed that templates are complementary to dfg in graphcodebert, one of existing structural-aware tasks.