SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In our experiments, for all the evaluated model configurations, we use nucleus sampling (Holtzman et al., 2020) with a top-p value of 0.95 to generate  $n = 6$  independent samples. For a budget of  $k \in [1, n]$  samples, we compute the aggregate score score@k (see Appendix D). On the discrete-valued metrics (CR and NIM), it is identical to pass@k,  $n$  (Chen et al., 2021), estimating the expected number of times the list of  $k$  candidates contains at least one successful compilation or match with ground-truth identifier. On the real-valued metrics (ISM and PM), it estimates the expectation of the maximum value of the corresponding metric given  $k$  chances.
Python Library for MGD. We are releasing an extensible Python library, multilspy, for interfacing between LMs and language servers using LSP. It can be used with multiple programming languages, static analyses and LMs. It also supports an approximate mechanism for MGD of black-box LMs with limited logit-masking support. Please refer to Appendix C for more details.
Table 1 shows the summary of the results for all of our experiments and metrics. For the "-MGD" configurations, we also report the relative improvement over the base model in parentheses, where the base model is the same model configuration without MGD. Below, we present a detailed evaluation.
4.1 Effect of MGD on Models across Parameter Scale and Architectures
We present results for all the models on Standard prompts described in Section 3.
Compilation Rate. As shown in Figure 2a, all the base models with MGD, including the smallest model CodeGen-350M for  $k \geq 4$ , outperform the largest model text-davinci-003, by maximum relative margin of  $16.55\%$  achieved by SantaCoder. All the models with MGD outperform their respective base models on Compilation Rate, by a relative improvement ranging between  $21.77\% - 24.69\%$ . TD-3-MGD outperforms TD-3 by a relative margin of  $18.52\%$ .
(a) CR
(b) NIM
(c) ISM
Figure 2: score@k for models with MGD and Standard prompt compared against base models. The values of  $k \in [1,6]$  are marked on the X-axis.
(d) PM
Next Identifier Match. As seen in Figure 2b, all the models with MGD outperform the respective base models, with a relative improvement of  $7.11\% - 8.92\%$ . The smallest model CodeGen-350M with MGD outperforms the much larger CodeGen-6B with a relative improvement of  $2.76\%$ . SantaCoder with MGD outperforms the larger CodeGen-6B by a relative margin of  $8.42\%$ , and even the largest model text-davinci-003 by  $2.60\%$ . TD-3-MGD outperforms TD-3 with a relative margin of  $5.81\%$ .
Identifier Sequence Match. Figure 2c shows that all the models with MGD outperform their respective base models on ISM, showing a relative improvement ranging between  $6.41\% -7.70\%$ . SantaCoder and CodeGen-2B with MGD outperform the larger CodeGen-6B with a relative margin of  $9.47\%$  and  $5.00\%$  respectively. TD-3-MGD outperforms TD-3 with a relative margin of  $5.24\%$
Prefix Match. Figure 2d shows percentage prefix match with ground truth. All the models with MGD outperform their respective base models with a relative improvement of  $5.89\% -6.82\%$ . Both SantaCoder and CodeGen-2B with MGD outperform the larger CodeGen-6B with a relative margin of  $8.08\%$  and  $4.30\%$ . TD-3-MGD outperforms TD-3 with a relative margin of  $3.03\%$ .
SC-MGD vs. TD-3. SC is a 1.1B parameter model whereas TD-3 has 175B parameters. Being a much larger model and possibly due to other differences in training, TD-3 does better than SC across all metrics. Interestingly, with MGD, SC-MGD outperforms TD-3 on CR (Figure 2a) and NIM (Figure 2b). ISM and PM are method-level metrics and the relative advantage of the larger model prevails. Even then, with a small increase in the sampling budget, from  $k = 1$  to  $k = 4$ , SC-MGD manages to surpass TD-3's performance with  $k = 1$  on ISM (Figure 2c) and PM (Figure 2d).
Summary. MGD improves the compileability of code significantly, across architectures and parameter scale, leading even the smallest CodeGen-350M with MGD to outperform the largest LM, text-davinci-003. We also see improvements in all ground truth agreement metrics. Notably, smaller LMs with MGD outperform larger LMs (CodeGen-350M with MGD outperforms text-davinci-003 in CR and NIM, CodeGen-2B with MGD outperforms CodeGen-6B on ISM and PM) across all metrics.
4.2 Effect of MGD and Prompt Augmentation Strategies
We choose the best-performing base model, SantaCoder, from Section 4.1 to study the effect of prompting when combined with MGD. Figure 3 shows results for SantaCoder-{Standard, classExprTypes, RLPG} and TD-3, compared against SantaCoder with respective prompts and MGD.
**Compilation Rate.** Figure 3a shows the results for Compilation Rate. We observe improvements in compilation rate with both the prompting techniques, classExprTypes and RLPG, with RLPG marginally outperforming classExprTypes. We note that SantaCoder with Standard prompt and MGD is able to relatively improve over both RLPG and classExprTypes augmentation by  $10.01\%$  and  $13.11\%$  respectively. Further, SantaCoder with RLPG and MGD is able to outperform SantaCoder-RLPG and SantaCoder-classExprTypes with a relative margin of  $17.70\%$  and  $21.02\%$  respectively, while increasing the margin of relative improvement over text-davinci-003 to  $24.70\%$ .
Next Identifier Match. As seen in Figure 3b, similar to compilation, both RLPG and classExprTypes prompt augmentation leads to improvement over the base model. However, SantaCoder with either prompt augmentations underperforms text-davinci-003. SantaCoder with MGD outperforms
(c) ISM
Figure 3: score@k for models with MGD and prompt augmentation compared against base models.
TD-3, and consequently, both SantaCoder-RLPG and SantaCoder-classExprTypes with a relative improvement of  $2.60\%$ ,  $3.51\%$  and  $4.14\%$  respectively. SantaCoder with prompting and MGD outperform their respective baselines (SantaCoder with prompting) by a relative margin in the range of  $5.23\% - 5.25\%$ . We note that SantaCoder with RLPG and MGD increases the relative improvement with respect to the largest model, text-davinci-003 to  $4.31\%$ .
Identifier Sequence Match. On ISM, SantaCoder with prompt augmentation and MGD is able to outperform its respective baseline by a relative improvement of  $4.78\% - 5.00\%$ , while both the prompt augmentations result in an improvement over the base model. SantaCoder with RLPG and MGD is able to significantly reduce the gap with text-davinci-003, underperforming it by just  $1.11\%$ .
Prefix Match. As seen in Figure 3d, SantaCoder with prompt augmentation and MGD is able to outperform its respective baseline by  $4.26\% - 4.87\%$ .
Summary. While prompt augmentation techniques help in improving performance on all metrics, we see further improvement with MGD augmentation, and conclude that the contributions by prompt augmentation and MGD are complementary. Notably, SantaCoder-RLPG with MGD improves the relative margin for compilation rate with respect to text-davinci-003 to  $24.70\%$  compared to the  $16.55\%$  improvement achieved by SC-MGD, or  $5.95\%$  improvement achieved by SC-RLPG.
4.3 Effect of MGD on Fill-in-the-middle (FIM) Decoding
Among the base models, SantaCoder supports the FIM modality. We evaluated SantaCoder with autoregressive and FIM decoding strategies and text-davinci-003, and compared them with respective configurations of SantaCoder with MGD. Similar to our observations with prompt augmentation, while FIM modality leads to improvements across all metrics, we see continued improvement when using both FIM and MGD. Due to space limitations, detailed results are in Appendix E. Motivated by the complementary nature of FIM and MGD, we further evaluated SC-FIM-classExprTypes-MGD, combining both prompt augmentation and FIM modality. Consistent with our findings, it leads to a further improvement over SC-FIM-classExprTypes, as seen in Figure 5 (Appendix E).
4.4 Effect of Identifier Complexity on Next Identifier Match
Identifier names in repositories can get specific and long (Karampatsis et al., 2020). Due to this, while commonly used APIs may get tokenized into single tokens, identifiers specific to the context of individual repositories, especially in private settings, can span multiple subtokens in the LM vocabulary. We define the complexity of an identifier as the mean number of subtokens required to decode it. We show that the ability of LMs to accurately predict the identifier name decreases sharply with an increase in identifier complexity, while augmenting them with MGD improves their performance. MGD provides an improvement in the range of  $21.79\% - 27.91\%$  compared to the base model without MGD, across parameter scales, for the highest identifier complexity, which is prevalent in more than  $36\%$  of the methods in DOTPROMPTS. The detailed results are available in Appendix F.
5 Generalizability Study
We curate MGDMicroBENCH as a micro benchmark (10 examples), spanning 3 programming languages (Java, C#, Rust), 4 coding scenarios, and requiring use of 2 different static analyses, to evaluate generalizability of MGD. The detailed results are discussed in Appendix H.
Different Programming Languages. MGD utilizes static analyses that help infer and enforce semantic constraints on the code under development. Such analyses are available through the Language Server Protocol (LSP) (lsp) for most programming languages, e.g., clangd for C/C++ (cla, 2020), Jedi for Python (Roeca, 2019) and Rust Analyzer for Rust (rus, 2018). The monitor used in MGD can be instantiated as a thin client around LSP. Supporting new languages is easy and doesn't necessitate changes to the monitor's static analysis interface. Hence, MGD is applicable to most programming languages. Results on MGDMICROBENCH demonstrate MGD for Java, C# and Rust.
Different Coding Scenarios. A monitor under MGD triggers when a specified precondition is met during code generation. Flexibility in defining preconditions allows MGD to be applied to a variety of code scenarios, as follows: 1) Instantiation of valid classes: A monitor triggers on 'new', invokes a static analysis that identifies instantiable classes from the local & global context to ensure only valid classes are instantiated. 2) switch over enum: A switch statement over enum uses named enums in case <val> to select branch in C# and Java. A monitor triggering on 'case' is used to generate valid enums. 3) Correct number of arguments: A stack-based monitor is implemented to guide for generation of right number of arguments to methods, that also handles nested function calls. 4) Joint monitoring for multiple properties: Unlike previous examples that monitored for a single static property, we instantiate 2 joint-monitors: a) jointly guiding for type-correct dereferences & along with right number of arguments, b) jointly guiding for valid switch enum branches & type-correct dereferences. MGDMicroBENCH (appendix H.2) provides results over all of these scenarios.
Different Static Analyses. MGD is able to utilize results from various static analyses to guide the generation of code with LMs. We demonstrate MGD over the following properties in MGDMICROBENCH: 1) Typestates (Strom & Yemini, 1986): Many APIs are stateful, and require callers to follow specific ordering of calls as part of the API contract. For example, a type representing a file handle would have a contract disallowing calls to read after close has been called. Such contracts can be expressed as finite state machines (FSMs), called typestates. 2) Session Types (Jespersen et al., 2015) ensure that messages between concurrent programs are sent and received in the expected order, following a protocol, and are specified as communicating FSMs. In MGDMICROBENCH, we demonstrate a monitor for Rust that utilizes typestate and session-type design analyses.
Limitations. Though static analysis is a mature field with strong theoretical foundations and several robust implementations of tools, analyzing partial and incomplete programs is still a difficult problem. In practice, editors such as Eclipse and Visual Studio support static analysis with heuristics. Though these heuristics are well-engineered and are widely used, they can be both imprecise (they can give incorrect suggestions) and incomplete (they can leave out correct suggestions). Satisfying functional-correctness specifications like pre/post-conditions and invariants is beyond the scope of this work. Consequently, though our results from guiding LMs using these analyses (through MGD) show improvements in quality metrics, additional steps such as testing and human inspection are needed to guarantee correctness of generated code.
Societal Impact. Software pervasively affects all aspects of our lives. With LMs being widely deployed as copilots and intelligent assistants to help developers write code, it is crucially important to develop tools like MGD to improve the quality of code generated by LMs (even if humans review and accept the suggestions given by LMs). Without such tools, we risk introducing bugs in code due to incorrect suggestions made by LMs, which has the potential to impact all of our lives negatively.
Amount of Compute. Our experiments do not involve any training, and we only perform inferences. We used machines of the following configurations: (1) CPU: 24-core AMD Epyc with 220GB RAM, GPU: Nvidia A100 80GB. (2) CPU: Intel Xeon(R) Platinum 8168 with 290GB RAM, GPU: Nvidia Tesla V100 16GB. For the experiments to evaluate text-davinci-003, we used the Azure API.
Pre-trained Models of Code. Many powerful pre-trained models have been designed for code. These include encoder-only models like CodeBERT (Feng et al., 2020), GraphCodeBERT Guo et al. (2020) and CuBERT (Kanade et al., 2020); encoder-decoder models like PLBART (Ahmad et al., 2021), CodeT5 (Wang et al., 2021) and AlphaCode (Li et al., 2022); or decoder-only models like Codex (Chen et al., 2021), GPT-J (Wang & Komatsuzaki, 2021), Austin et al. (2021), GPT-Neo (Black et al., 2021), GPT-NeoX (Black et al., 2022), CodeParrot (Tunstall et al., 2022), PolyCoder (Xu et al., 2022), Incoder (Fried et al., 2022), CodeGen (Nijkamp et al., 2023), SantaCoder (Allal et al., 2023) and StarCoder (Li et al., 2023); or unified models like UniXCoder (Guo et al., 2022). Our monitor-guided decoding works only with logits and hence can be used with any model.
Global Context of Code. Hellendoorn & Devanbu (2017) build an n-gram model with a cache to track directory-level context. Xu et al. (2021) use locality based on directory-structure in retrieval-augmented modeling (Khandelwal et al., 2019). Many approaches use static analysis or previous generations (Zhang et al., 2023) to extract relevant code. They use the relevant context to either augment the prompt (Shrivastava et al., 2022; Pei et al., 2023; Zhang et al., 2023) or embeddings (Pashakhanloo et al., 2022; Ding et al., 2022) presented as input to the LM. Due to the limit on the prompt or embedding size, these approaches filter information through random walks (Pashakhanloo et al., 2022), classification (Shrivastava et al., 2022) or using fixed pruning strategies (Ding et al., 2022).
As we neither augment the prompt nor use extra embeddings, we do not need to prune the global context. We let the static analysis generate completion suggestions using the entire repository-level context. Zan et al. (2022) and Zhou et al. (2022) retrieve information from library documentation for prompt augmentation. Our static analysis analyzes libraries along with the repository-level source code. Several of these techniques require architecture modifications (Pashakhanloo et al., 2022; Ding et al., 2022) or finetuning (Zan et al., 2022; Ding et al., 2022; Pei et al., 2023). We use a simple interface between logits and static analysis with a frozen LM. Most of these approaches, excluding (Xu et al., 2021; Zhang et al., 2023), use one-time a priori retrieval. In contrast, we provide token-level guidance by invoking static analysis on demand. Our method is complementary to all the above approaches as they all try to condition the generation by modifying the input to the LM whereas we apply output-side constraints by reshaping the logits.
Syntactic and Semantic Constraints. There are two primary lines of work to enforce syntactic and semantic constraints on code generation, based on specialized modeling and through constrained decoding. GNN2NAG (Brockschmidt et al., 2019) and NSG (Mukherjee et al., 2021) are examples of the first and use attribute grammars. They are respectively evaluated on expressions that do not use user-defined methods or on methods with class-level context. We consider repository context for method-level completion. Unlike these approaches, our work is applicable to off-the-shelf LMs. PICARD (Scholak et al., 2021) and Synchronesh (Poesia et al., 2022) are constrained decoding approaches similar to ours. They use incremental parsing for syntactic validity and design domain-specific semantic checks to ensure semantic validity. Both are evaluated on SQL, and Synchronesh additionally considers domain-specific languages for visualization and calendar applications. In comparison, we target generation of general-purpose programming languages with focus on semantic constraints like type-consistency, API protocols, etc. using static analysis over repository context.
In this work, we show how to use repository-wide information computed by static analysis (specifically, type-based analysis) using a stateful monitor as an interface, to improve quality of code generated by LMs. Our experimental results show the potential for significant quality improvements for code generation using this approach. Our approach is complementary to prompt augmentation techniques. It allows smaller models to achieve better or competitive performance compared to much larger models. This could open up the possibility of using smaller models directly within IDEs, alongside our monitor, as an alternative to the use of remotely-hosted Large LMs (LLMs), reducing inference costs and improving privacy. Our method is general and applicable to various coding scenarios where LMs are used generatively, such as code refactoring, code repair, or code completion, even if the repositories are in a transient state. We plan to expand the scope of MGD to more languages and deeper semantic analyses such as pre/post-conditions for which advanced constrained decoding methods (including backtracking and beam-search) might be needed.
We thank the authors RLPG and of the LMs (CG, SC, TD-3) used in our work. We are grateful to Sheng Chen and Shi Chen from Microsoft, and the author of OLSP, Predrag Nikolic, for their generous help in answering our queries about interfacing with LSP. Our paper also benefited significantly from the constructive feedback by the reviewers.
Official page for Language Server Protocol. URL https://microsoft.github.io/language-server-protocol/.
Protocol Buffers. https://protobuf.dev/, 2008. Accessed: May 14, 2023.
Project Lombok. https://projectlombok.org/, 2009. Accessed: May 14, 2023.
Eclipse JDT LS, September 2016. URL https://projects.eclipse.org/projects/eclipse.jdt.1s.
rust-analyzer, 2018. URL https://rust-analyzer.github.io/.
What is clangd?, 2020. URL https://clangd.llvm.org/.
Language Server Implementations. https://microsoft.github.io/language-server-protocol/implementors/servers/, 2023. Accessed: May 14, 2023.
Ahmad, W. U., Chakraborty, S., Ray, B., and Chang, K.-W. Unified pre-training for program understanding and generation, 2021.
Allal, L. B., Li, R., Kocetkov, D., Mou, C., Akiki, C., Ferrandis, C. M., Muennighoff, N., Mishra, M., Gu, A., Dey, M., and others. SantaCoder: don't reach for the stars! arXiv preprint arXiv:2301.03988, 2023.
Austin, J., Odena, A., Nye, M., Bosma, M., Michalewski, H., Dohan, D., Jiang, E., Cai, C., Terry, M., Le, Q., and Sutton, C. Program Synthesis with Large Language Models, August 2021. URL http://arxiv.org/abs/2108.07732. arXiv:2108.07732 [cs].
Bavarian, M., Jun, H., Tezak, N., Schulman, J., McLeavey, C., Tworek, J., and Chen, M. Efficient training of language models to fill in the middle, 2022.
Black, S., Gao, L., Wang, P., Leahy, C., and Biderman, S. Gpt-neo: Large scale autoregressive language modeling with mesh-tensorflow. If you use this software, please cite it using these metadata, 58, 2021.
Black, S., Biderman, S., Hallahan, E., Anthony, Q., Gao, L., Golding, L., He, H., Leahy, C., McDonell, K., Phang, J., and others. Gpt-neox-20b: An open-source autoregressive language model. arXiv preprint arXiv:2204.06745, 2022.
Brockschmidt, M., Allamanis, M., Gaunt, A. L., and Polozov, O. Generative Code Modeling with Graphs. In International Conference on Learning Representations, 2019. URL https://openreview.net/forum?id=Bke4KsA5FX.
Chen, M., Tworek, J., Jun, H., Yuan, Q., Pinto, H. P. d. O., Kaplan, J., Edwards, H., Burda, Y., Joseph, N., Brockman, G., and others. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374, 2021.
Crichton, W.Typed design patterns for the functional era. In Proceedings of the 1st ACM SIGPLAN International Workshop on Functional Software Architecture, FUNARCH 2023, pp. 40-48, New York, NY, USA, 2023. Association for Computing Machinery. ISBN 9798400702976. doi: 10.1145/3609025.3609477. URL https://doi.org/10.1145/3609025.3609477.
Crichton, W., Navarro, J., Semeniuta, D., Zeng, S., and Gupta, S. Session types, November 2019. URL https://stanford-cs242.github.io/f19/lectures/09-1-session-types.html.
Appendix: Monitor-Guided Decoding of Code LMs with Static Analysis of Repository Context
- Monitor for the Running Example
- Data Set Curation Details
- Experimental Setup - Additional Details
- Calculation of score@k
- Effect of MGD on Fill-in-the-middle (FIM) Decoding - Complete Results
- Effect of Identifier Complexity on Next Identifier Match - Complete Results
- Impact of MGD on Inference Time
- Generalizability aspects of MGD
- CodeQL Query for Identifying Evaluation Target Methods
Examples of code generation with MGD
A Monitor for the Running Example
Figure 4: Monitor to guide generation of type-consistent identifiers for the code in Figure 1.
Figure 4 shows how the monitor interacts with the LM decoder for the example in Figure 1. Initially, the monitor,  $M_{\varphi}$  is in the wait state  $s_0$ . Given the code completion input,  $x_1, x_2, \dots, x_n$ ,  $M_{\varphi}$  first evaluates  $pre(s_0, x_1, \dots, x_n)$ . Since  $x_n =$  (')' (the dot symbol indicating object dereferencing in Java),  $pre(s_0, x_1, \dots, x_n)$  evaluates to true, and subsequently, in accordance with Eq. (4), the static analysis  $A_{\varphi}$  is invoked, which determines the input prompt to be in accordance with the property  $\varphi$ , and resolves the type for the completion point to be ServerNode. Builder, as shown in the annotated AST in Figure 1(b).  $A_{\varphi}$  then returns the set of identifiers consistent with the resolved type - {withIp, withPort, newServerNode, ...}, transitioning the monitor  $M_{\varphi}$  to state  $s_1$ .  $M_{\varphi}$  then calculates  $m = \text{maskgen}(s_1, V)$ , which masks out, for example, the token host (as inferred by SantaCoder in Figure 1(b)). Concurrently, the input is tokenized and the LM  $L_{\theta}$  provides inferred logits,  $\ell$  for the next token. The output logits  $\ell$  from  $L_{\theta}$  and mask  $m$  are combined as  $\ell \oplus m$  to obtain the modified logits, which is then softened and a token is sampled -with in this case. The monitor then invokes update  $(s_1, \text{with})$  to transition to  $s_2$ . Note that with the state transition, newServerNode is pruned from the set of identifiers, as the sampled token with does not prefix it.  $L_{\theta}$  provides logits for the