Towards Compositional Generalization in Code Search

Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 10743 - 10750 December 7-11, 2022 ©2022 Association for Computational Linguistics Towards Compositional Generalization in Code Search Hojae Han♠∗, Seung-won Hwang♠†, Shuai Lu⋄, Nan Duan⋄, Seungtaek Choi♣ ♠Seoul National University, {stovecat, seungwonh}@snu.ac.kr ⋄Microsoft Research Asia, {shuailu, nanduan}@microsoft.com ♣Riiid, seungtaek.choi@riiid.co Abstract We study compositional generalization, which aims to generalize on unseen combinations of seen structural elements, for code search. Unlike existing approaches of partially pursu- ing this goal, we study how to extract struc- tural 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 em- pirically 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 en- hancing structural context around variables. 1 Introduction We study compositional generalization in the con- text of code search, which aims to generalize on unseen combinations of previously seen elements, or building blocks (Qiu et al., 2022). Following the categorization in the NLP domain (Kim and Linzen, 2020), 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 combina- tion 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 ob- served while its combination was not (structural generalization). Existing code search models such as Code- BERT (Feng et al., 2020) regarding code as an unstructured sequence (Figure 2a) often fail to gen- eralize when the lexical perturbation without chang- ing structures and labels is applied in test codes, known as AdvTest evaluation (Lu et al., 2021). To- wards better generalization, SynCoBERT (Wang ∗Work done during internship at Microsoft. †Corresponding author. 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 k > i: print(i) … Seen (a) Lexical generalization. ... try: a /= b except ZeroDivisionError: a = 0 ... try: _ except _: raise _ while _: _ if _: break … try: while True: l = foo(l) if len(l) == 0: break except TypeError: raise bar() … … while i > 0: i = i-k if func(i): break … (b) Structural generalization. Figure 1: Motivation. Templates enhance compositional generalization between seen and unseen codes. et al., 2021b) extracts an abstract syntax tree (AST; Figure 2b) to incorporate syntactic structure into the model. On the other hand, GraphCode- BERT (Guo et al., 2021) extracts a data flow graph (DFG; Figure 2c), to represent variable relations within code. From structure-aware pretraining ob- jectives, these two models significantly enhance their accuracy in code search, both in original and adversarial settings. Meanwhile, utilizing compos- able building blocks, like the colored …

Content truncated to 500 words