\section{Background}
\vspace{-0.2cm}

\subsection{Semantic Code Search}
\vspace{-0.2cm}
Semantic code search (SCS) is the process of retrieving a relevant code snippet based on a textual description of its functionality, also referred to as \textit{query}.
Let $\mathcal{C}$ be a database of code snippets $\mathbf{c}^i$. For each $\mathbf{c}^i \in \mathcal{C}$, there is a textual description of its functionality $\mathbf{q}^i$. In the example in Figure~\ref{fig:overview}, the query $\mathbf{q}^i$ is ``\textit{Load all tables from dataset}''.
Let $r$ be an indicator function such that $r(\mathbf{q}^i, \mathbf{c}^j) = 1$ if $i=j$; and 0 otherwise.
Given some query $\mathbf{q}$ the goal of SCS is to find $\mathbf{c}^*$ such that $r(\mathbf{q}, \mathbf{c}^*) = 1$. We assume that for each $\mathbf{q}^*$ there is exactly one such $\mathbf{c}^*$.\footnote{This is not the case in CoSQA dataset. For the sake of consistency, we perform the evaluation repeatedly, leaving only one correct code snippet among the candidates at a time, while removing the others.}
Here we look to construct a model which takes as input a pair of query and a candidate code snippet: $(\mathbf{q}^i, \mathbf{c}^j)$ and assign the pair a probability $\hat{r}^{ij}$ for being a correct match. Following the common practice in information retrieval, we evaluate the performance of the model based on how high the correct answer $\mathbf{c}^*$ is ranked among a number of incorrect, or distractor instances $\{\mathbf{c}\}$. This set of distractor instances can be the entire codebase $\mathcal{C}$, or a subset of the codebase obtained through heuristic filtering, or another ranking method. 

\vspace{-0.1cm}
\subsection{Neural Models for Semantic Code Search}
\vspace{-0.1cm}

Past works handling programs and code have focused on enriching their models with incorporating more semantic and syntactic information from code~\citep{DBLP:conf/acl/AhmadCRC20,DBLP:conf/acl/ChoiBNL21,DBLP:conf/emnlp/Shi0D0HZS21,DBLP:conf/iclr/ZugnerKCLG21}. 
Some prior works have cast the SCS as a sequence classification task, where the code is represented as a textual sequence and input pair $(\mathbf{q}^i, \mathbf{c}^j)$ is concatenated with a special separator symbol into a single sequence, and the output is the score $\hat{r}^{ij}$: $\hat{r}^{ij} = f(\mathbf{q}^i, \mathbf{c}^j)$. The function $f$ performing the classification can be any sequence classification model, e.g. BERT~\citep{DBLP:conf/naacl/DevlinCLT19}. 

Alternatively, one can define separate networks for independently representing the query ($f$), the code ($g$) and measuring the similarity between them: $\hat{r}^{ij} = sim(f(\mathbf{q}^i), g(\mathbf{c}^j))$. This allows one to design the code encoding network $g$ with additional program-specific information, such as abstract syntax trees~\citep{DBLP:journals/pacmpl/AlonZLY19, DBLP:conf/icse/ZhangWZ0WL19} or control flow graphs~\citep{DBLP:journals/nn/GuLGWZXL21, DBLP:conf/sigsoft/ZhaoH18}. Separating two modalities of natural language and code also allows further enrichment of code representation by adding contrastive learning objectives~\citep{DBLP:conf/emnlp/0001JZA0S21, DBLP:conf/sigir/BuiYJ21}. In these approaches, the original code snippet $\mathbf{c}$ is automatically modified with semantic-preserving transformations, such as variable renaming, to introduce versions of the code snippet - $\mathbf{c}^\prime$ with the exact same functionality. Code encoder $g$ is then trained with an appropriate contrastive loss, such as Noise Contrastive Estimation (NCE)~\citep{DBLP:journals/jmlr/GutmannH10}, or InfoNCE~\citep{DBLP:journals/corr/abs-1807-03748}. 



\vspace{-0.0cm}
\textbf{Limitations}
However, there is also merit in reviewing how we represent and use the textual \textit{query} to help guide the SCS process.
Firstly, existing work derive a single embedding for the entire query. The problem of representing longer text inputs is being actively researched in the community and remains an open problem~\citep{DBLP:journals/corr/abs-2004-05150,DBLP:conf/rep4nlp/BhathenaWD20, DBLP:journals/corr/abs-1904-00157}. This means that specific details or nested subqueries of the query may be omitted or not represented properly - getting lost in the embedding. Secondly, prior approaches make the decision after a single pass over the code snippet. 
This ignores cases where reasoning about a query requires multiple steps and thus - multiple look-ups over the code, as is for example in cases with nested subqueries.
Our proposed approach - $NS^3$ - attempts to address these issues by breaking down the query into smaller phrases based on its semantic parse and locating each of them in the code snippet. This should allow us to match compositional and longer queries to code more precisely.
\section{Neural Modular Code Search}
\label{sec:model}
\vspace{-0.2cm}

We propose to supplement the query with a loose structure resembling its semantic parse, as illustrated in Figure~\ref{fig:overview}. We follow the parse structure to break down the query into smaller, semantically coherent parts, so that each corresponds to an individual execution step.
The steps are taken in succession by a neural module network composed from a layout that is determined from the semantic parse of the query (Sec.~\ref{sec:layout}).
The neural module network is composed by stacking ``modules'', or jointly trained networks, of distinct types, each carrying out a different functionality.

\vspace{-0.2cm}
\paragraph{Method Overview}
In this work, we define two types of neural modules - \textit{entity discovery} module (denoted by $E$; Sec.~\ref{sec:entity-module}) and \textit{action} module (denoted by $A$; Sec~\ref{sec:action-module}). The entity discovery module estimates semantic relatedness of each code token $c^j_i$ in the code snippet \begin{small}$\mathbf{c}^j=[c^j_1, \ldots, c^j_N]$\end{small} to an entity mentioned in the query -- e.g. ``\textit{all tables}'' or ``\textit{dataset}'' as in  Figure~\ref{fig:overview}.
The action module estimates the likelihood of each code token to be related to an (unseen) entity affected by the action in the query e.g. ``\textit{dataset}'' and ``\textit{load from}'' correspondingly, conditioned on the rest of the input (seen), e.g. ``\textit{all tables}''. The similarity of the predictions of the entity discovery and action modules measures how well the code matches that part of the query.  
The modules are nested - the action modules are taking as input part of the output of another module - and the order of nesting is decided by the semantic parse layout. In the rest of the paper we refer to the inputs of a module as its \textit{arguments}. 

Every input instance fed to the model is a 3-tuple $(\mathbf{q}^i, s_{q^i}, \mathbf{c}^j)$ consisting of a natural language query $\mathbf{q}^i$, the query's semantic parse $s_{q^i}$, a candidate code (sequence) $\mathbf{c}^j$. The goal is producing a binary label $\hat{r}^{ij}=1$ if the code is a match for the query, and 0 otherwise.
The layout of the neural module network, denoted by $L(s_{q^i})$, is created from the semantic structure of the query $s_{q^i}$. 
During inference, given $(\mathbf{q}^i, s_{q^i}, \mathbf{c}^j)$ as input the model instantiates a network based on the layout, passes $\mathbf{q}^i$, $\mathbf{c}^j$ and $s_{q^i}$ as inputs, and obtains the model prediction $\hat{r}^{ij}$. This pipeline is illustrated in Figure~\ref{fig:overview}, and details about creating the layout of the neural module network are presented in Section~\ref{sec:layout}.

During training, we first perform noisy supervision pretraining for both modules. Next, we perform end-to-end training, where in addition to the query, its parse, and a code snippet, the model is also provided a gold output label $r(\mathbf{q}^i, \mathbf{c}^j)=1$ if the code is a match for the query, and $r(\mathbf{q}^i, \mathbf{c}^j)=0$ otherwise. These labels provide signal for joint fine-tuning of both modules (Section~\ref{sec:training}).


\vspace{-0.1cm}
\subsection{Module Network Layout}\label{sec:layout}
\vspace{-0.1cm}
Here we present our definition of the structural representation $s_{q^i}$ for a query $\mathbf{q}^i$, and introduce how this structural representation is used for dynamically constructing the neural module network, i.e. building its layout $L(s_{q^i})$.

\vspace{-0.2cm}
\paragraph{Query Parsing}
To infer the representation $s_{q^i}$, we pair the query (\textit{e.g.}, ``\textit{Load all tables from dataset}'', as in Figure~\ref{fig:overview}), with a simple semantic parse that looks similar to: \begin{small}\texttt{DO WHAT [ (to/from/in/\ldots) WHAT, WHEN, WHERE, HOW, etc]}\end{small}. 
Following this semantic parse, we break down the query into shorter semantic phrases using the roles of different parts of speech. Nouns and noun phrases correspond to data entities in code, and verbs describe actions or transformations performed on the data entities. Thus, data and transformations are separated and handled by separate neural modules -- an \textit{entity discovery module} $E$ and an \textit{action module} $A$. 
We use a Combinatory Categorial Grammar-based (CCG) semantic parser~\citep{DBLP:journals/corr/abs-1207-1420,DBLP:conf/emnlp/ArtziLZ15} to infer the  semantic parse  $s_{q^i}$ for the natural language query $\mathbf{q}^i$. Parsing is described in further detail in Section~\ref{sec:parser} and Appendix~\ref{sec:app-parsing}.



\vspace{-0.2cm}
\paragraph{Specifying Network Layout}
In the layout \begin{small}$L(s_{q^i})$\end{small}, every noun phrase (e.g., ``\textit{dataset}" in Figure~\ref{fig:overview}) will be passed through the entity discovery module $E$. Module $E$ then produces a probability score $e_k$ for every token $c^j_k$ in the code snippet $\mathbf{c}^j$ to indicate its semantic relatedness to the noun phrase:
\begin{small}
$E(\text{``\textit{dataset}''}, \mathbf{c}^j) = [e_1, e_2, \ldots, e_N].$
\end{small}
Each verb in $s_{q^i}$ (e.g., ``\textit{load}'' in Figure~\ref{fig:overview}) will be passed through an action module:
\begin{small}$A(\text{``\textit{load}''}, \mathbf{p}^i, \mathbf{c}^j) = [a_1, a_2, \ldots, a_N]$\end{small}.
Here, $\mathbf{p}^i$ is the span of arguments to the verb (action) in query $\mathbf{q}^i$, consisting of children of the verb in the parse $s_{q^i}$ (e.g. subject and object arguments to the predicate \textit{``load''}); $a_1, \ldots, a_N$ are estimates of the token scores $e_1, \ldots, e_N$ for an entity from $\mathbf{p}^i$.
The top-level of the semantic parse is always an action module. Figure~\ref{fig:overview} also illustrates preposition \textsc{FROM} used with ``\textit{dataset}'', handling which is described in Section~\ref{sec:action-module}.

\vspace{-0.1cm}
\subsection{Entity Discovery Module}
\label{sec:entity-module}
\vspace{-0.1cm}


The entity discovery module receives a string that references a data entity. Its goal is to identify tokens in the code that have high relevance to that string. 
The architecture of the module is shown in Figure~\ref{fig:entitydiscovery}. Given an entity string, ``\textit{dataset}'' in the example, and a sequence of code tokens $[c^j_1, \ldots, c^j_N]$, entity module first obtains contextual code token representation using RoBERTa model that is initialized from CodeBERT-base checkpoint. The resulting embedding is passed through a two-layer MLP to obtain a score for every individual code token $c^j_k$ : $0 \leq e_k \leq 1$. Thus, the total output of the module is a vector of scores: $[e_1, e_2, \ldots, e_N]$. 
To prime the entity discovery module for measuring relevancy between code tokens and input, we fine-tune it with noisy supervision, as detailed below. 

\vspace{-0.2cm}
\paragraph{Noisy Supervision}
We create noisy supervision for the entity discovery module by using keyword matching and a Python static code analyzer.
For the keyword matching, if a code token is an exact match for one or more tokens in the input string, its supervision label is set to 1, otherwise it is 0. Same is true if the code token is a substring or a superstring of one or more input string tokens. For some common nouns we include their synonyms (e.g. ``\textit{map}'' for ``\textit{dict}''), the full list of those and further details are presented in Appendix~\ref{sec:app-entity-module}.

\begin{wrapfigure}{r}{0.45\textwidth}
\centering
\vspace{-0.3cm}
\includegraphics[width=0.45\textwidth]{images/entity-module}
\vspace{-0.5cm}
\caption{\small Entity module architecture.}
\label{fig:entitydiscovery}
\vspace{-0.4cm}
\end{wrapfigure}

We used the static code analyzer to extract information about statically known data types. We cross-matched this information with the query to discover whether the query references any datatypes found in the code snippet. If that is the case, the corresponding code tokens are assigned supervision label 1, and all the other tokens are assigned to $0$. 
In the pretraining we learned on equal numbers of $(query, code)$ pairs from the dataset, as well as randomly mismatched pairs of queries and code snippets to avoid creating bias in the entity discovery module. 

\vspace{-0.2cm}
\subsection{Action Module}
 \label{sec:action-module}
\vspace{-0.1cm}

First, we discuss the case where the action module has only entity module inputs.
Figure~\ref{fig:action-estimation} provides a high-level illustration of the action module. 
In the example, for the query ``\emph{Load all tables from dataset}'', the action module receives only part of the full query -- ``\emph{Load all tables from ???}''. 
Action module then outputs token scores for the masked argument -- ``\emph{dataset}''. 
If the code snippet corresponds to the query, then the action module should be able to deduce this missing part from the code and the rest of the query. For consistency, we always mask the last data entity argument. 
We pre-train the action module using the output scores of the entity discovery module as supervision.

\begin{wrapfigure}{r}{0.47\textwidth}
\vspace{-0.5cm}
\centering
\includegraphics[width=0.47\textwidth]{images/action-module}
\vspace{-0.5cm}
\caption{\small Action module architecture.}
\label{fig:action-estimation}
\vspace{-0.3cm}
\end{wrapfigure}
More specifically, after masking one input, we compute joint embeddings for verb and preposition for each data entity argument. This is done with a 2-layer MLP model, as illustrated in the left-most part of Figure~\ref{fig:action-estimation}. 
If a data entity does not have a preposition associated with it, that part of the input is filled with zeros. The joint verb-preposition embedding is stacked with the code token embedding $t_k$ and entity discovery module output for that token, this is referenced in the middle part of Figure~\ref{fig:action-estimation}. This vector is passed through a transformer encoder model, followed by a 2-layer MLP and a sigmoid layer to output token score $a_k$, illustrated in the right-most part of the Figure~\ref{fig:action-estimation}.

\vspace{-0.1cm}
\subsection{Model Prediction}
\vspace{-0.2cm}

The final score $\hat{r}^{ij}=f(\mathbf{q}^i, \mathbf{c}^j)$ is computed based on the similarity of action and entity discovery module output scores. Formally, for an action module with verb $x$ and parameters $\mathbf{p}^x = [p^x_1, \ldots, p^x_k]$, the final model prediction is the dot product of respective outputs of action and entity discovery modules: $\hat{r}^{ij} = A(x, p^x_1, \ldots, p^x_{k-1}) \cdot E(p^x_k)$.
Since the action module estimates token scores for the entity affected by the verb, if its prediction is far from the truth - then either the action is not found in the code, or it is not fully corresponding to the query, for example, in the code snippet tables are loaded from web, instead of a dataset.
We normalize this score to make it a probability. If this is the only action in the query, this probability score will be the output of the entire model for ($\mathbf{q}^i, \mathbf{c}^j$)  pair: $\hat{r}^{ij}$, otherwise $\hat{r}^{ij}$ will be the product of probability scores of all nested actions in the layout.


\vspace{-0.2cm}
\paragraph{Compositional query with nested actions} 
Consider a compositional query ``\textit{Load all tables from dataset using \emph{Lib} library}''. Here action with verb ``\textit{Load from}'' has an additional argument ``\textit{using}'' -- also an action -- with an entity argument ``\textit{\emph{Lib} library}''. In case of nested actions, 
we flatten the layout by taking the conjunction of individual action similarity scores. 
Formally, for two verbs $x$ and $y$ and their corresponding arguments $\mathbf{p}^x = [p^x_1, \ldots, p^x_k]$ and $\mathbf{p}^y = [p^y_1, \ldots, p^y_l]$ in a layout that looks like: $A(x, \mathbf{p}^x, A(y, \mathbf{p}^y))$, 
the output of the model is the conjunction of similarity scores computed for individual action modules:
$sim(A(x, p^x_1, \ldots, p^x_{k-1}), E(p^x_k)) \cdot sim(A(y, p^y_1, \ldots, p^y_{l-1}), E(p^y_l))$. This process is repeated until all remaining $\mathbf{p}^x$ and $\mathbf{p}^y$ are data entities. 
This design ensures that code snippet is ranked highly if both actions are ranked highly, we leave explorations of alternative handling approaches for nested actions to future work.

\vspace{-0.2cm}
\subsection{Module Pretraining and Joint Fine-tuning}
\label{sec:training}
\vspace{-0.1cm}
We train our model through supervised pre-training, as is discussed in Sections~\ref{sec:entity-module} and~\ref{sec:action-module}, followed by end-to-end training. End-to-end training objective is binary classification - given a pair of query $\mathbf{q}^i$ and code $\mathbf{c}^j$, the model predicts probability $\hat{r}^{ij}$ that they are related. In the end-to-end training, we use positive examples taken directly from the dataset - ($\mathbf{q}^i$, $\mathbf{c}^i$), as well as negative examples composed through the combination of randomly mismatched queries and code snippets. The goal of end-to-end training is fine-tuning parameters of entity discovery and action modules, including the weights of the RoBERTA models used for code token representation. 

Batching is hard to achieve for our model, so for the interest of time efficiency we do not perform inference on all distractor code snippets in the code dataset. Instead, for a given query we re-rank top-K highest ranked code snippets as outputted by some baseline model, in our evaluations we used CodeBERT. Essentially, we use our model in a re-ranking setup, this is common in information retrieval and is known as L2 ranking. 
We interpret the probabilities outputted by the model as ranking scores. More details about this procedure are provided in Section~\ref{sec:eval-metrics}.
\section{Related work}
\vspace{-0.2cm}
Different deep learning models have proved quite efficient when applying to programming languages and code. Prior works have studied and reviewed the uses of deep learning for code analysis in general and code search in particular~\citep{DBLP:journals/corr/abs-2101-11149, DBLP:conf/nips/LuGRHSBCDJ21}.

A number of approaches to deep code search is based on creating a relevance-predicting model between text and code. 
\citep{DBLP:conf/icse/GuZ018} propose using RNNs for embedding both code and text to the same latent space. 
On the other hand, \citep{DBLP:journals/tkdd/LingWWPMXLWJ21} capitalizes the inherent graph-like structure of programs to formulate code search as graph matching.
A few works propose enriching the models handling code embedding by adding additional code analysis information, such as semantic and dependency parses~\citep{DBLP:conf/cikm/DuSWSHZ21, DBLP:conf/msr/AkbarK19}, variable renaming and statement permutation~
\citep{DBLP:journals/corr/abs-2012-01028}, as well as structures such as abstract syntax tree of the program~\citep{DBLP:conf/acl/HaldarWXH20, DBLP:conf/kbse/WanSSXZ0Y19}.
A few other approaches have dual formulations of code retrieval and code summarization~\citep{DBLP:conf/kbse/ChenZ18,  DBLP:conf/www/YaoPS19, DBLP:conf/www/YeXZHWZ20,DBLP:conf/sigir/BuiYJ21}
In their work \citep{DBLP:journals/corr/abs-2008-12193} propose considering the code search scenario where short annotative descriptions of code snippets are provided. Appendix~\ref{sec:app-related-work} discusses more related work.

\section{Conclusion}
\vspace{-0.2cm}
We presented NS$^3$ a symbolic method for semantic code search based on neural module networks. 
Our method represents the query and code in terms of actions and data entities.
In contrast to existing code search methods, NS$^3$ better captures the compositional nature of queries. In an extensive evaluation, we show that this method works better than strong but unstructured baselines. 
We further study model's generalization capacities, robustness, and sensibility of outputs in a series of additional experiments.





\section*{References}


\section*{Checklist}




\begin{enumerate}


\item For all authors...
\begin{enumerate}
  \item Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
    \answerYes{}
  \item Did you describe the limitations of your work?
    \answerYes{See the Discussion paragraph under Section~\ref{sec:analysis-and-case-study}}
  \item Did you discuss any potential negative societal impacts of your work?
    \answerNA{}
  \item Have you read the ethics review guidelines and ensured that your paper conforms to them?
    \answerYes{}
\end{enumerate}


\item If you are including theoretical results...
\begin{enumerate}
  \item Did you state the full set of assumptions of all theoretical results?
    \answerNA{}
        \item Did you include complete proofs of all theoretical results?
    \answerNA{}
\end{enumerate}


\item If you ran experiments...
\begin{enumerate}
  \item Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
    \answerYes{Code and data in supplemental material}
  \item Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
    \answerYes{See paragraph titled Hyperparameters in Section~\ref{sec:experiment-setting}}
        \item Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
    \answerNo{}
        \item Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
    \answerYes{}
\end{enumerate}


\item If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
\begin{enumerate}
  \item If your work uses existing assets, did you cite the creators?
    \answerYes{}
  \item Did you mention the license of the assets?
    \answerNo{}
  \item Did you include any new assets either in the supplemental material or as a URL?
    \answerNo{}
  \item Did you discuss whether and how consent was obtained from people whose data you're using/curating?
    \answerNA{}
  \item Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
    \answerNA{}
\end{enumerate}


\item If you used crowdsourcing or conducted research with human subjects...
\begin{enumerate}
  \item Did you include the full text of instructions given to participants and screenshots, if applicable?
    \answerNA{}
  \item Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
    \answerNA{}
  \item Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
    \answerNA{}
\end{enumerate}


\end{enumerate}





\section{Experiments}
\vspace{-0.2cm}

\subsection{Experiment Setting}
\label{sec:experiment-setting}
\vspace{-0.1cm}


\paragraph{Dataset}
We conduct experiments on two datasets: Python portion of the \textbf{CodeSearchNet} (CSN) \citep{DBLP:journals/corr/abs-1909-09436}, and \textbf{CoSQA} \citep{DBLP:conf/acl/HuangTSGXJZD21}. We parse all queries with the CCG parser, Section~\ref{sec:parser}, excluding unparsable examples from further experiments. This leaves us with approximately 40\% of the CSN dataset and 70\% of the CoSQA dataset, the exact data statistics are available in Appendix~\ref{sec:app-experiment-settings} in Table~\ref{tab:data-size}. We believe, that the difference in success rate of the parser between the two datasets can be attributed to the fact that CSN dataset, unlike CoSQA, does not contain real code search queries, but rather consists of docstrings, which are used as approximate queries. More details and examples can be found in Appendix~\ref{sec:app-docstrings}.
For our baselines, we use the parsed portion of the dataset for fine-tuning to make the comparison fair. In addition, we also experiment with fine-tuning all models on an even smaller subset of CodeSearchNet dataset, using only 5K and 10K examples for fine-tuning. The goal is testing whether modular design makes NS3 more sample-efficient.

All experiment and ablation results discussed in Sections~\ref{sec:results},\ref{sec:ablation-studies} and \ref{sec:analysis-and-case-study} are obtained on the test set of CSN for models trained on CSN training data, or WebQueryTest \citep{DBLP:conf/nips/LuGRHSBCDJ21} --  a small natural language web query dataset of document-code pairs -- for models trained on CoSQA dataset.

\vspace{-0.2cm}
\paragraph{Evaluation and Metrics}\label{sec:eval-metrics}
We follow CodeSearchNet's original approach for evaluation for a test instance $(q, c)$, comparing the output against outputs over a fixed set of 999 distractor code snippets.
We use two evaluation metrics: Mean Reciprocal Rank (MRR) and Precision@K (P@K) for K=1, 3, and 5, see Appendix~\ref{sec:app-metrics} for definitions and further details.

Following a common approach in information retrieval, we perform two-step evaluation. In the first step, we obtain CodeBERT's output against 999 distractors. In the second step, we use $NS^3$ to re-rank the top 10 predictions of CodeBERT. This way the evaluation is much faster, since unlike our modular approach, CodeBERT can be fed examples in batches. And as we will see from the results, we see improvement in final performance in all scenarios. 


\vspace{-0.2cm}
\paragraph{Compared Methods}
We compare $\textsc{NS}^3$ with various state-of-the-art methods, including some traditional approaches for document retrieval and pretrained large NLP language models. (1) \textbf{BM25} is a ranking method to estimate the relevance of documents to a given query. 
(2) \textbf{RoBERTa (code)} is a variant of RoBERTa \citep{DBLP:journals/corr/abs-1907-11692} pretrained on the CodeSearchNet corpus.
(3) \textbf{CuBERT} \citep{DBLP:conf/icml/KanadeMBS20} is a BERT Large model pretrained on 7.4M Python files from GitHub. 
(4) \textbf{CodeBERT} is an encoder-only Transformer model trained on unlabeled source code via masked language modeling (MLM) and replaced token detection objectives. 
(5) \textbf{GraphCodeBERT} is a pretrained Transformer model using MLM, data flow edge prediction, and variable alignment between code and the data flow. 
(6) \textbf{GraphCodeBERT*} is a re-ranking baseline. We used the same setup as for NS3, but used GraphCodeBERT to re-rank the top-10 predictions of the CodeBERT model. 

\vspace{-0.2cm}
\paragraph{Query Parser}\label{sec:parser}
We started by building a vocabulary of predicates for common action verbs and entity nouns, such as ``\textit{convert}'', ``\textit{find}'', ``\textit{dict}'', ``\textit{map}'', etc.
For those we constructed the lexicon (rules) of the parser. We have also included ``catch-all'' rules, for parsing sentences with less-common words. To increase the ratio of the parsed data, we preprocessed the queries by removing preceding question words, punctuation marks, etc. Full implementation of our parser including the entire lexicon and vocabulary can be found at \url{https://anonymous.4open.science/r/ccg_parser-4BC6}. Details are available in Appendix~\ref{sec:app-parsing}.

\vspace{-0.2cm}
\paragraph{Pretrained Models}
Action and entity discovery modules each embed code tokens with a RoBERTa model, that has been initialized from a checkpoint of pretrained CodeBERT model \footnote{\url{https://huggingface.co/microsoft/codebert-base}}. We fine-tune these models during the pretraining phases, as well as during final end-to-end training phase. 


\begin{wraptable}{r}{0.43\textwidth}
\vspace{-0.4cm}
    \hspace{-0.1cm}
    \scalebox{0.76}{
        \begin{tabular}{lllll} \\
        \toprule
        \multirow{2}{*}{Method} & \multicolumn{4}{c}{CoSQA} \\ 
        & MRR & P@1 & P@3 & P@5 \\
        \midrule
        BM25 & 0.103 & 0.05 & 0.119 & 0.142\\
       
        RoBERTa (code) & 0.279  & 0.159 & 0.343 & 0.434 \\
        CuBERT & 0.127 & 0.067 & 0.136 & 0.187 \\
        CodeBERT & 0.345 & 0.175 &  0.42 & 0.54 \\
        GraphCodeBERT & 0.435 & 0.257 & 0.538 & 0.628\\
        GraphCodeBERT* & 0.462 & 0.314 & 0.547 & 0.632\\
        \midrule
        NS$^3$  & \textbf{0.551} & \textbf{0.445} & \textbf{0.619} & \textbf{0.668}\\
        \midrule
        Upper-bound & 0.736 & 0.724 & 0.724 & 0.724\\
        \bottomrule
        \end{tabular}
    }
    \caption{\small Mean Reciprocal Rank(MRR) and Precision@1/@3/@5 (higher is better) for different methods trained on CoSQA dataset.}
     \label{tab:cosqa}
    \vspace{-0.2cm}
\end{wraptable}

\begin{table*}[t]
\vspace{-0.7cm}
    \centering
    \scalebox{0.78}{
        \begin{tabular}{lllllllllllll} \\
        \toprule
            \multirow{2}{*}{Method} & \multicolumn{4}{c}{CSN} & \multicolumn{4}{c}{CSN-10K} & \multicolumn{4}{c}{CSN-5K} \\ 
             & MRR & P@1 & P@3 & P@5 & MRR & P@1 & P@3 & P@5 & MRR & P@1 & P@3 & P@5\\ 
            \midrule
            BM25 & 0.209 & 0.144 & 0.230 & 0.273 & 0.209 & 0.144 & 0.230 & 0.273 & 0.209 & 0.144 & 0.230 & 0.273 \\
           
            RoBERTa (code) & 0.842 & 0.768 & 0.905 & 0.933 & 0.461 & 0.296 & 0.545 & 0.664 & 0.290 & 0.146 & 0.324 & 0.438\\
            CuBERT & 0.225 & 0.168  & 0.253 & 0.294 & 0.144 & 0.081 & 0.166 & 0.214 & 0.081 & 0.030 & 0.078 & 0.118 \\
            CodeBERT & 0.873 & 0.803 & 0.939 & 0.958 & 0.69 & 0.55 & 0.799 & 0.873 & 0.680 & 0.535 & 0.794 & 0.870\\
            GraphCodeBERT & 0.812 & 0.725 & 0.880 & 0.919 & 0.786 & 0.684 & 0.859 & 0.901 & 0.773 & 0.677 & 0.852 & 0.892\\
            GraphCodeBERT* & 0.883 & 0.820 & 0.941 & 0.962 & 0.780 & 0.683 & 0.858 & 0.904 & 0.765 & 0.662 & 0.846 & 0.894\\
            \midrule
            NS$^3$ & \textbf{0.924} & \textbf{0.884} & \textbf{0.959} & \textbf{0.969} & \textbf{0.826} & \textbf{0.753} & \textbf{0.886} & \textbf{0.908} & \textbf{0.823}    & \textbf{0.751} & \textbf{0.881} & \textbf{0.913} \\
            \midrule
            Upper-bound & 0.979 &  &  &  & 0.939 &  &  &  & 0.936 &   &   &  \\
            \bottomrule
        \end{tabular}
    }
    \caption{Mean Reciprocal Rank (MRR) and Precision@1/@3/@5 (higher is better) for methods trained on different subsets from CodeSearchNet dataset.}
    \label{tab:results-csn}
    \vspace{-0.1cm}
\end{table*}

\vspace{-0.2cm}
\paragraph{Hyperparameters}
The MLPs in entity discovery and action modules have 2 layers with input dimension of 768. We use dropout in these networks with rate $0.1$. The learning rate for pretraining and end-to-end training phases was chosen from the range of 1e-6 to 6e-5. We use early stopping with evaluation on unseen validation set for model selection during action module pretraining and end-to-end training. For entity discovery model selection we performed manual inspection of produced scores on unseen examples. For fine-tuning the CuBERT, CodeBERT and GraphCodeBERT baselines we use the hyperparameters reported in their original papers. For RoBERTa (code), we perform the search for learning rate during fine-tuning stage in the same interval as for our model. For model selection on baselines we also use early stopping.


\vspace{-0.1cm}
\subsection{Results}\label{sec:results}
\vspace{-0.2cm}

\paragraph{Performance Comparison}
Tables~\ref{tab:results-csn}~and~\ref{tab:cosqa} present the performance evaluated on testing portion of CodeSearchNet dataset, and WebQueryTest dataset correspondingly. As it can be seen, our proposed model outperforms the baselines. 

Our evaluation strategy improves performance only if the correct code snippet was ranked among the top-10 results returned by the CodeBERT model, so rows labelled ``Upper-bound'' report best possible performance with this evaluation strategy. 
 
\begin{figure*}[t]
\begin{subfigure}[h]{0.45\textwidth}
\centering
\includegraphics[width=\textwidth]{images/barchart-table6}
\caption{}
\label{fig:max-depth}
\end{subfigure}
\begin{subfigure}[h]{0.52\textwidth}
\centering
\includegraphics[width=\textwidth]{images/barchart-table7}
\caption{}
\end{subfigure}
\vspace{-0.2cm}
\caption{\small We report Precision@1 scores. (a) Performance of our proposed method and baselines broken down by average number of arguments per action in a single query. (b) Performance of our proposed method and baselines broken down by number of arguments in queries with a single action.}
\label{fig:arguments-number}
\vspace{-0.1cm}
\end{figure*}

\paragraph{Query Complexity vs. Performance}
Here we present the breakdown of the performance for our method vs baselines, using two proxies for the complexity and compositionality of the query. The first one is the maximum depth of the query. We define the maximum depth as the maximum number of nested action modules in the query. The results for this experiment are presented in Figure~\ref{fig:max-depth}. As we can see, $NS^3$ improves over the baseline in all scenarios. It is interesting to note, that while CodeBERT achieves the best performance on queries with depth 3+, our model's performance peaks at depth = 1. We hypothesize that this can be related to the automated parsing procedure, as parsing errors are more likely to be propagated in deeper queries. Further studies with carefully curated manual parses are necessary to better understand this phenomenon. 

\begin{figure*}[t]
\centering
\begin{subfigure}[h]{0.31\textwidth}
\centering
\includegraphics[width=\textwidth]{images/training-ablations.pdf}
\caption{Effect of training}
\label{fig:training-ablations}
\end{subfigure}
\begin{subfigure}[h]{0.33\textwidth}
\centering
\includegraphics[width=\textwidth]{images/norm-ablations.pdf}
\caption{Score normalization}
\label{fig:norm-ablations}
\end{subfigure}
\begin{subfigure}[h]{0.32\textwidth}
\centering
\includegraphics[width=\textwidth]{images/similarity-ablations}
\caption{Similarity measure}
\label{fig:similarity-ablations}
\end{subfigure}
\vspace{-0.1cm}
\caption{\small Performance of $NS^3$ on the test portion of CSN dataset with different ablation variants. (a) Skipping one, or both pretraining procedures, and only training end-to-end. (b) Using no normalization on output scores (\textbf{None}), action-only or entity discovery-only, and both. (c) Performance with different options for computing action and entity discovery output similarities.}
\label{fig:ablations}
\vspace{-0.4cm}
\end{figure*}

Another proxy for the query complexity we consider, is the number of data arguments to a single action module. While the previous scenario is breaking down the performance by the depth of the query, here we consider its ``width''. We measure the average number of entity arguments per action module in the query. In the parsed portion of our dataset we have queries that range from 1 to 3 textual arguments per action verb. The results for this evaluation are presented in Figure~\ref{fig:arguments-number}. As it can be seen, there is no significant difference in performances between the two groups of queries in either CodeBERT or our proposed method - $NS^3$.

\vspace{-0.1cm}
\subsection{Ablation Studies}\label{sec:ablation-studies}
\vspace{-0.2cm}
\paragraph{Effect of Pretraining}
In an attempt to better understand the individual effect of the two modules as well as the roles of their pretraining and training procedures, we performed two additional ablation studies. In the first one, we compare the final performance of the original model with two versions where we skipped part of the pretraining. The model noted as ($NS^3-AP$) was trained with pretrained entity discovery module, but no pretraining was done for action module, instead we proceeded to the end-to-end training directly. For the model called $NS^3-(AP\&EP)$, we skipped both pretrainings of the entity and action modules, and just performed end-to-end training. Figure~\ref{fig:training-ablations} demonstrates that combined pretraining is important for the final performance. 
Additionally, we wanted to measure how effective the setup was without end-to-end training. The results are reported in Figure~\ref{fig:training-ablations} under the name $NS^3-E2E$. There is a huge performance dip in this scenario, and while the performance is better than random, it is obvious that end-to-end training is crucial for $NS^3$.

\vspace{-0.3cm}
\subparagraph{Score Normalization}
We wanted to determine the importance of output normalization for the modules to a proper probability distribution. In Figure~\ref{fig:norm-ablations} we demonstrate the performance achieved using no normalization at all, normalizing either action or entity discovery module, or normalizing both. In all cases we used L1 normalization, since our output scores are non-negative. The version that is not normalized at all performs the worst on both datasets. The performances of the other three versions are close on both datasets.

\vspace{-0.3cm}
\subparagraph{Similarity Metric}
Additionally, we experimented with replacing the dot product similarity with a different similarity metric. In particular, in Figure~\ref{fig:similarity-ablations} we compare the performance achieved using dot product similarity, L2 distance, and weighted cosine similarity. The difference in performance among different versions is marginal.

\vspace{-0.1cm}
\subsection{Analysis and Case Study}\label{sec:analysis-and-case-study}
\vspace{-0.2cm}

Appendix~\ref{sec:app-additional-analysis} contains additional studies on model generalization, such as handling completely unseen actions and entities, as well as the impact of the frequency of observing an action or entity during training has on model performance. 


\vspace{-0.4cm} 
\subparagraph{Case Study}
Finally, we demonstrate some examples of the scores produced by our modules at different stages of training. Figure~\ref{fig:case-study1} shows module score outputs for two different queries and with their corresponding code snippets. The first column shows the output of the entity discovery module after pretraining, while the second and third columns demonstrate the outputs of entity discovery and action modules after the end-to-end training. We can see that in the first column the model identifies syntactic matches, such as ``folder'' and a list comprehension, which ``elements'' could be related too. After fine-tuning we can see there is a wider range of both syntactic and some semantic matches present, e.g. ``dirlist'' and ``filelist'' are correctly identified as related to ``folders''.

\vspace{-0.4cm}
\paragraph{Perturbed Query Evaluation}
\label{sec:perturbed-query}
In this section we study how sensitive the models are to small changes in the query $\mathbf{q}^i$, so that it no longer correctly describes its corresponding code snippet $\mathbf{c}^i$. 
Our expectation is that evaluating a sensitive model on $\mathbf{c}^i$ will rate the original query higher than the perturbed one. Whereas a model that tends to over-generalize and ignore details of the query will likely rate the perturbed query similar to the original. We start from 100 different pairs $(\mathbf{q}^i, \mathbf{c}^i)$, that both our model and CodeBERT predict correctly. 
\begin{wrapfigure}{r}{0.42\textwidth}
\vspace{-0.4cm}
\centering
\includegraphics[width=0.42\textwidth]{images/barchart-table5}
\caption{\small Ratio of the perturbed query score to the original query score  (lower is better) on CSN dataset.}
\label{fig:perturbed_query_eval}
\vspace{-0.3cm}
\end{wrapfigure}
We limited our study to queries with a single verb and a single data entity argument to that verb. For each pair we generated perturbations of two kinds, with 20 perturbed versions for every query. For the first type of perturbations, we replaced query's data argument with a data argument sampled randomly from another query. For the second type, we replaced the verb argument with another randomly sampled verb.
To account for calibration of the models, we measure the change in performance through ratio of the perturbed query score over original query score (lower is better). 
The results are shown in Figure~\ref{fig:perturbed_query_eval}, labelled ``$V(arg_1) \rightarrow  V(arg_2)$'' and ``$V_1(arg) \rightarrow V_2(arg)$''.


\begin{figure}[t]
\centering
\vspace{-0.7cm}
\includegraphics[width=0.92\textwidth]{images/case-study.png}
\vspace{-0.1cm}
\caption{\small 
Token scores outputted by the modules at different stages of training. Darker highlighting means higher score. The leftmost and middle columns show output scores of the entity discovery module after pretraining, and the end-to-end training correspondingly. The rightmost column shows the scores of the action module after the end-to-end training. }
\label{fig:case-study1}
\vspace{-0.2cm}
\end{figure}

\vspace{-0.3cm}
\paragraph{Discussion}
One of the main requirements for the application of our proposed method is being able to construct a semantic parse of the retrieval query. In general, it is reasonable to expect the users of the SCS to be able to come up with a formal representation of the query, e.g. by representing it in a form similar to SQL or CodeQL. However, due to the lack of such data for training and testing purposes, we implemented our own parser, which understandably does not have perfect performance since we are dealing with open-ended sentences. 


\section{Introduction}
\vspace{-0.2cm}

\begin{wrapfigure}{r}{0.45\textwidth} 
\centering
\vspace{-0.8cm}
\includegraphics[width=0.44\textwidth]{images/motivating-figure}
\vspace{-0.3cm}
\caption{\small \textbf{Motivating Example}. To match query ``\textit{Navigate folders}'' on a code snippet, we find all references (token spans) to entity ``\textit{folders}'' in code (e.g., paths and directories) using various linguistic cues (Step 1). Then we look for cues in code that indicate the identified instances of ``\textit{folders}" are being iterated through -- i.e., ``\textit{navigate}" (Step 2).
}
\label{fig:motivating}
\vspace{-0.3cm}
\end{wrapfigure}

The increasing scale of software repositories makes retrieving relevant code snippets more challenging. Traditionally, source code retrieval has been limited to keyword~\citep{DBLP:conf/icsm/Reiss09, DBLP:conf/wcre/LuSWLD15} or regex~\citep{DBLP:conf/wcre/BullTMG02} search. Both rely on the user providing
the exact keywords appearing in or around the sought code.
However, neural models enabled new approaches for retrieving code from a textual description of its functionality, a task known as \textit{semantic code search (SCS)}.
A model like Transformer~\citep{DBLP:conf/nips/VaswaniSPUJGKP17} can map a database of code snippets and natural language queries to a shared high-dimensional space. Relevant code snippets are then retrieved by searching over this embedding space using a predefined similarity metric, or a learned distance function~\citep{DBLP:conf/icml/KanadeMBS20, DBLP:conf/emnlp/FengGTDFGS0LJZ20, DBLP:conf/cikm/DuSWSHZ21}.
Some of the recent works capitalize on the rich structure of the code, and employ graph neural networks for the task~ \citep{DBLP:conf/iclr/GuoRLFT0ZDSFTDC21, DBLP:journals/corr/abs-2111-02671}. 

Despite impressive results on SCS, current neural approaches are far from satisfactory in dealing with a wide range of natural-language queries, especially on ones with compositional language structure.
Encoding longer text into a dense vector is an open problem for neural language models, as neural networks are not believed to be extracting systematic rules from data~\citep{ DBLP:journals/corr/abs-2004-05150,DBLP:conf/rep4nlp/BhathenaWD20, DBLP:journals/corr/abs-1904-00157}.
Not only does this a) affect the performance, but it can b) drastically reduce a model's value for the users, because compositional queries such as ``\emph{Check that directory does not exist before creating it}'' require performing multi-step reasoning on code.

We suggest overcoming these challenges by introducing a modular workflow based on the semantic structure of the query. 
Our approach is based on the intuition of how an engineer would approach a SCS task. For example, in performing search for code that navigates folders in Python they would first only pay attention to code that has cues about operating with paths, directories or folders. Afterwards, they would seek indications of iterating through some of the found objects or other entities in the code related to them. In other words, they would perform multiple steps of different nature - i.e. finding indications of specific types of data entities, or specific operations. Figure~\ref{fig:motivating} illustrates which parts of the code would be important to indicate that they have found the desired code snippet at each step. We attempt to imitate this process in this work. To formalize the decomposition of the query into such steps, we take inspiration from the idea that code is comprised of data, or entities, and transformations, or actions, over data. Thus, a SCS query is also likely to describe the code in terms of data entities and actions.

\begin{figure}[t]
\vspace{-0.8cm}
\centering
\includegraphics[width=0.99\textwidth]{images/overview.pdf}
\vspace{-0.1cm}
\caption{\small \textbf{Overview of the NS3 approach}. We illustrate the pipeline of processing for an example query ``\textit{Load all tables from dataset}''. Parsed query is used for deciding the positions of entity discovery and action modules in the neural module network layout. Each entity discovery module receives a noun/noun phrase as input, and outputs relatedness scores for code tokens, which are passed as input to an action module. Action module gets scores for all its children in the parse-tree, except one, which is masked, and the goal is predicting, cloze-style, what are the relatedness scores for the missing argument.}
\vspace{-0.2cm}
\label{fig:overview}
\end{figure}

We break down the task of matching the query into smaller tasks of matching individual data entities and actions. In particular, we aim to identify parts of the code that indicate the presence of the corresponding data or action. We tackle each part with a distinct type of network -- a neural module. 
Using the semantic parse of the query, we construct the layout of how modules' outputs should be linked according to the relationships between data entities and actions. Correspondingly, this layout specifies how the modules should be combined into a single neural module network (NMN)~\citep{DBLP:conf/cvpr/AndreasRDK16}. NMN is ``evaluated'' on the candidate code snippet to decide whether the query and code are related. This pipeline is outlined in Figure~\ref{fig:overview}, and described in further detail in Section~\ref{sec:model}.


This approach has the following advantages. First, semantic parse captures the compositionality of a query. Second, it mitigates the challenges of encoding a long text by removing the need for it, as each of the modules handle only a small portion of the query. A smaller portion of the query is less likely to require compositional reasoning from an individual neural module. Finally, applying the neural modules in a succession can potentially mimic staged reasoning necessary for SCS. 

We evaluate our proposed NS$^3$ model on two SCS datasets - CodeSearchNet (CSN) \citep{DBLP:journals/corr/abs-1909-09436} and CoSQA/WebQueryTest \citep{DBLP:conf/acl/HuangTSGXJZD21}. Additionally, we experiment with a limited training set size of CSN of 10K and 5K examples. We find that NS$^3$ provides large improvements upon baselines in all cases. To measure how well the model captures compositional properties of the queries, we experiment with performing adversarial modifications to the query and showing that compared to the baselines, NS$^3$ is more likely to correctly recognize modified queries. We also research the effect on queries that require multi-staged reasoning over code by breaking down the performance by query's compositional complexity, and demonstrate that our proposed method handles this task better than the baselines.

Our main contributions are:  
(i) We propose looking at SCS as a compositional task that requires multi-step reasoning.
(ii) We present an implementation of the aforementioned paradigm based on NMNs. 
(iii) We demonstrate that our proposed model provides a large improvement on a number of well-established baseline models.
(iv) We perform additional studies to evaluate the capacity of our model to handle compositional queries. 


\section{Experiment Settings}\label{sec:app-experiment-settings}
\begin{table}[t]
    \centering
    \resizebox{0.5\textwidth}{!}{%
        \begin{tabular}{lrrr} \\
        \toprule
        Dataset & Train & Valid & Test\\ 
        \midrule
        CodeSearchNet & 162801 & 8841 & 8905 \\
        CoSQA & 14210 & - & - \\
        WebQueryTest & - & - & 662 \\
    \bottomrule
    \end{tabular}
    }
    \caption{Dataset statistics after parsing.} 
    \label{tab:data-size}
\end{table}

\subsection{Evaluation Metrics}\label{sec:app-metrics}
(1) \textbf{MRR} evaluates a list of code snippets.
The reciprocal rank for MRR is computed as $\frac{1}{rank}$, where $rank$ is the position of the correct code snippet when all code snippets are ordered by their predicted similarity to the sample query. 
(2) \textbf{P@K}
is the proportion of the top-\textit{K} correct snippets closest to the given query. 
For each query, if the correct code snippet is among the first \textit{K} retrieved code snippets P@K=1, otherwise it is 0.


\subsection{Parsing}\label{sec:app-parsing}
We build on top of the NLTK Python package for our implementation of the CCG parser. In attempt to parse as much of the datasets as possible, we preprocessed the queries by removing preceding question words (e.g. ``\textit{How to}''), punctuation marks, and some specific words and phrases, e.g. those that specify a programming language or version, such as ``\textit{in Python}'' and ``\textit{Python 2.7}''. For a number of entries in CSN dataset which only consisted of a noun or a noun phrase, we appended a Load verb to make it a valid sentence, assuming that it was implied, so that, for example, ``\textit{video page}'' became ``\textit{Load video page}''. This had the adverse effect in cases of noisy examples, where the docstring did not specify the intention or functionality of the function, and only said ``\textit{wrapper}'', for example.
          
\subsection{Failed parses}\label{sec:app-docstrings}
As mentioned before, we have encountered many noisy examples and here provide samples of such examples that could not be parsed. These include cases where the docstring contains URLs, is not in English, consists of multiple sentences, or has code in it, which is often either signature of the function, or a usage example. Specific samples of queries that we couldn't parse are included in Table~\ref{tab:failed-parses}.

\begin{table}[h]
\begin{tabular}{ll}
\toprule
& Example not parsed \\
URL & From http://cdn37.atwikiimg.com/sitescript/pub/dksitescript/FC2.site.js\\
\midrule
Signature & \begin{tabular}[c]{@{}l@{}}:param media\_id:\\ :param self: bot\\ :param text: text of message\\ :param user\_ids: list of user\_ids for creating group or one user\_id for send to one person    \\ :param thread\_id: thread\_id\end{tabular} \\
\midrule
Multi-sentence & \begin{tabular}[c]{@{}l@{}}Assumed called on Travis, to prepare a package to be deployed\\ This method prints on stdout for Travis. \\ Return is obj to pass to sys.exit() directly\end{tabular} \\
\midrule
Noisy & bandwidths are inaccurate, as we don't account for parallel transfers here \\
\bottomrule
\end{tabular}
\caption{Example queries that were not included due to query parsing errors}
\label{tab:failed-parses}
\end{table}
\vspace{-0.2cm}

\section{Entity Module}\label{sec:app-entity-module}

\section{Additional Experiments}\label{sec:app-additional-analysis}
\subsection{Unseen Entities and Actions}
We wanted to see how well different models adapt to new entities and actions that were not seen during training. For that end we measured the performance of the models when broken down on queries with a different number of unseen entities (from 0 to 3+) and action (0 and 1). The results are presented in Figure~\ref{fig:unseen}. It can be seen that NS3 is very sensitive to unseen terms, whereas CodeBERT's performance stays the same. 

\begin{figure}[h]
\begin{subfigure}[h]{0.45\textwidth}
\centering
\includegraphics[width=\textwidth]{images/appendix_unseen_actions.pdf}
\caption{Unseen Actions}
\end{subfigure}
\begin{subfigure}[h]{0.45\textwidth}
\centering
\includegraphics[width=\textwidth]{images/appendix_unseen_entities.pdf}
\caption{Unseen Entities}
\end{subfigure}
\vspace{-0.2cm}
\caption{Performance of CodeBERT and NS3 models when broken down by the number of unseen entities or actions in the test queries. Evaluated on CSN test set.}
\label{fig:unseen}
\end{figure}

\subsection{Times an Entity or an Action Was Seen}
In addition to the last experiment, we wanted to measure the performance broken down by how many times an entity or an action verb was seen during the training. The results of this experiment are reported in Figure~\ref{fig:times-seen}.
For the breakdown by the number of times an action was seen, the performance almost follows a bell curve. The performance increases with verbs that were seen only a few times. On the other hand, very frequent actions are probably too generic and not specific enough (e.g. load and get). For the entities we see that the performance is only affected when none of the entities in the query has been seen. This is understandable, as in these cases an action modules don’t get any information to go by, so the result is also bad. CodeBERT model in both scenarios has more or less the same performance independently of the number of times an action or an entity was seen. 
\begin{figure}[h]
\begin{subfigure}[h]{0.45\textwidth}
\centering
\includegraphics[width=\textwidth]{images/appendix_breakdown_by_entity.png}
\caption{Entities}
\end{subfigure}
\begin{subfigure}[h]{0.45\textwidth}
\centering
\includegraphics[width=\textwidth]{images/appendix_breakdown_by_action.png}
\caption{Actions}
\end{subfigure}
\vspace{-0.2cm}
\caption{Performance of CodeBERT and NS3 models when broken down by the number of times an entity or an action was seen during the training. Evaluated on CSN test set.}
\label{fig:times-seen}
\end{figure}

\section{Additional Examples}\label{sec:app-case-study}
Figure~\ref{fig:additional-case-studies} contains more illustrations of the output scores of the action and entity discovery modules captured at different stages of training. The queries shown here are the same, but this time they are evaluated on different functions.

\begin{figure}[h]
\vspace{-0.2cm}
\centering
\includegraphics[width=0.9\textwidth]{images/case-study-appendixes.png}
\caption{The leftmost column shows output scores of the entity discovery module after pretraining for the \textbf{\textit{entity}} of the query. The middle column shows the scores after completing the end-to-end training. The rightmost column shows the scores of the action module. Darker highlighting demonstrates higher score.}
\label{fig:additional-case-studies}
\vspace{-0.2cm}
\end{figure}


\section{Related Work}\label{sec:app-related-work}









\section{Experiment Settings}\label{sec:app-experiment-settings}
\begin{table}[t]
    \centering
    \resizebox{0.5\textwidth}{!}{%
        \begin{tabular}{lrrr} \\
        \toprule
        Dataset & Train & Valid & Test\\ 
        \midrule
        CodeSearchNet & 162801 & 8841 & 8905 \\
        CoSQA & 14210 & - & - \\
        WebQueryTest & - & - & 662 \\
    \bottomrule
    \end{tabular}
    }
    \caption{Dataset statistics after parsing.} 
    \label{tab:data-size}
\end{table}

\subsection{Evaluation Metrics}\label{sec:app-metrics}
(1) \textbf{MRR} evaluates a list of code snippets.
The reciprocal rank for MRR is computed as $\frac{1}{rank}$, where $rank$ is the position of the correct code snippet when all code snippets are ordered by their predicted similarity to the sample query. 
(2) \textbf{P@K}
is the proportion of the top-\textit{K} correct snippets closest to the given query. 
For each query, if the correct code snippet is among the first \textit{K} retrieved code snippets P@K=1, otherwise it is 0.


\subsection{Parsing}\label{sec:app-parsing}
We build on top of the NLTK Python package for our implementation of the CCG parser. In attempt to parse as much of the datasets as possible, we preprocessed the queries by removing preceding question words (e.g. ``\textit{How to}''), punctuation marks, and some specific words and phrases, e.g. those that specify a programming language or version, such as ``\textit{in Python}'' and ``\textit{Python 2.7}''. For a number of entries in CSN dataset which only consisted of a noun or a noun phrase, we appended a Load verb to make it a valid sentence, assuming that it was implied, so that, for example, ``\textit{video page}'' became ``\textit{Load video page}''. This had the adverse effect in cases of noisy examples, where the docstring did not specify the intention or functionality of the function, and only said ``\textit{wrapper}'', for example.
          
\subsection{Failed parses}\label{sec:app-docstrings}
As mentioned before, we have encountered many noisy examples and here provide samples of such examples that could not be parsed. These include cases where the docstring contains URLs, is not in English, consists of multiple sentences, or has code in it, which is often either signature of the function, or a usage example. Specific samples of queries that we couldn't parse are included in Table~\ref{tab:failed-parses}.

\begin{table}[h]
\begin{tabular}{ll}
\toprule
& Example not parsed \\
URL & From http://cdn37.atwikiimg.com/sitescript/pub/dksitescript/FC2.site.js\\
\midrule
Signature & \begin{tabular}[c]{@{}l@{}}:param media\_id:\\ :param self: bot\\ :param text: text of message\\ :param user\_ids: list of user\_ids for creating group or one user\_id for send to one person    \\ :param thread\_id: thread\_id\end{tabular} \\
\midrule
Multi-sentence & \begin{tabular}[c]{@{}l@{}}Assumed called on Travis, to prepare a package to be deployed\\ This method prints on stdout for Travis. \\ Return is obj to pass to sys.exit() directly\end{tabular} \\
\midrule
Noisy & bandwidths are inaccurate, as we don't account for parallel transfers here \\
\bottomrule
\end{tabular}
\caption{Example queries that were not included due to query parsing errors}
\label{tab:failed-parses}
\end{table}
\vspace{-0.2cm}

\section{Entity Module}\label{sec:app-entity-module}

\section{Additional Experiments}\label{sec:app-additional-analysis}
\subsection{Unseen Entities and Actions}
We wanted to see how well different models adapt to new entities and actions that were not seen during training. For that end we measured the performance of the models when broken down on queries with a different number of unseen entities (from 0 to 3+) and action (0 and 1). The results are presented in Figure~\ref{fig:unseen}. It can be seen that NS3 is very sensitive to unseen terms, whereas CodeBERT's performance stays the same. 

\begin{figure}[h]
\begin{subfigure}[h]{0.45\textwidth}
\centering
\includegraphics[width=\textwidth]{images/appendix_unseen_actions.pdf}
\caption{Unseen Actions}
\end{subfigure}
\begin{subfigure}[h]{0.45\textwidth}
\centering
\includegraphics[width=\textwidth]{images/appendix_unseen_entities.pdf}
\caption{Unseen Entities}
\end{subfigure}
\vspace{-0.2cm}
\caption{Performance of CodeBERT and NS3 models when broken down by the number of unseen entities or actions in the test queries. Evaluated on CSN test set.}
\label{fig:unseen}
\end{figure}

\subsection{Times an Entity or an Action Was Seen}
In addition to the last experiment, we wanted to measure the performance broken down by how many times an entity or an action verb was seen during the training. The results of this experiment are reported in Figure~\ref{fig:times-seen}.
For the breakdown by the number of times an action was seen, the performance almost follows a bell curve. The performance increases with verbs that were seen only a few times. On the other hand, very frequent actions are probably too generic and not specific enough (e.g. load and get). For the entities we see that the performance is only affected when none of the entities in the query has been seen. This is understandable, as in these cases an action modules don’t get any information to go by, so the result is also bad. CodeBERT model in both scenarios has more or less the same performance independently of the number of times an action or an entity was seen. 
\begin{figure}[h]
\begin{subfigure}[h]{0.45\textwidth}
\centering
\includegraphics[width=\textwidth]{images/appendix_breakdown_by_entity.png}
\caption{Entities}
\end{subfigure}
\begin{subfigure}[h]{0.45\textwidth}
\centering
\includegraphics[width=\textwidth]{images/appendix_breakdown_by_action.png}
\caption{Actions}
\end{subfigure}
\vspace{-0.2cm}
\caption{Performance of CodeBERT and NS3 models when broken down by the number of times an entity or an action was seen during the training. Evaluated on CSN test set.}
\label{fig:times-seen}
\end{figure}

\section{Additional Examples}\label{sec:app-case-study}
Figure~\ref{fig:additional-case-studies} contains more illustrations of the output scores of the action and entity discovery modules captured at different stages of training. The queries shown here are the same, but this time they are evaluated on different functions.

\begin{figure}[h]
\vspace{-0.2cm}
\centering
\includegraphics[width=0.9\textwidth]{images/case-study-appendixes.png}
\caption{The leftmost column shows output scores of the entity discovery module after pretraining for the \textbf{\textit{entity}} of the query. The middle column shows the scores after completing the end-to-end training. The rightmost column shows the scores of the action module. Darker highlighting demonstrates higher score.}
\label{fig:additional-case-studies}
\vspace{-0.2cm}
\end{figure}


\section{Related Work}\label{sec:app-related-work}









\section{Conclusion}
\vspace{-0.2cm}
We presented NS$^3$ a symbolic method for semantic code search based on neural module networks. 
Our method represents the query and code in terms of actions and data entities.
In contrast to existing code search methods, NS$^3$ better captures the compositional nature of queries. In an extensive evaluation, we show that this method works better than strong but unstructured baselines. 
We further study model's generalization capacities, robustness, and sensibility of outputs in a series of additional experiments.





\section{Introduction}
\vspace{-0.2cm}

\begin{wrapfigure}{r}{0.45\textwidth} 
\centering
\vspace{-0.8cm}
\includegraphics[width=0.44\textwidth]{images/motivating-figure}
\vspace{-0.3cm}
\caption{\small \textbf{Motivating Example}. To match query ``\textit{Navigate folders}'' on a code snippet, we find all references (token spans) to entity ``\textit{folders}'' in code (e.g., paths and directories) using various linguistic cues (Step 1). Then we look for cues in code that indicate the identified instances of ``\textit{folders}" are being iterated through -- i.e., ``\textit{navigate}" (Step 2).
}
\label{fig:motivating}
\vspace{-0.3cm}
\end{wrapfigure}

The increasing scale of software repositories makes retrieving relevant code snippets more challenging. Traditionally, source code retrieval has been limited to keyword~\citep{DBLP:conf/icsm/Reiss09, DBLP:conf/wcre/LuSWLD15} or regex~\citep{DBLP:conf/wcre/BullTMG02} search. Both rely on the user providing
the exact keywords appearing in or around the sought code.
However, neural models enabled new approaches for retrieving code from a textual description of its functionality, a task known as \textit{semantic code search (SCS)}.
A model like Transformer~\citep{DBLP:conf/nips/VaswaniSPUJGKP17} can map a database of code snippets and natural language queries to a shared high-dimensional space. Relevant code snippets are then retrieved by searching over this embedding space using a predefined similarity metric, or a learned distance function~\citep{DBLP:conf/icml/KanadeMBS20, DBLP:conf/emnlp/FengGTDFGS0LJZ20, DBLP:conf/cikm/DuSWSHZ21}.
Some of the recent works capitalize on the rich structure of the code, and employ graph neural networks for the task~ \citep{DBLP:conf/iclr/GuoRLFT0ZDSFTDC21, DBLP:journals/corr/abs-2111-02671}. 

Despite impressive results on SCS, current neural approaches are far from satisfactory in dealing with a wide range of natural-language queries, especially on ones with compositional language structure.
Encoding longer text into a dense vector is an open problem for neural language models, as neural networks are not believed to be extracting systematic rules from data~\citep{ DBLP:journals/corr/abs-2004-05150,DBLP:conf/rep4nlp/BhathenaWD20, DBLP:journals/corr/abs-1904-00157}.
Not only does this a) affect the performance, but it can b) drastically reduce a model's value for the users, because compositional queries such as ``\emph{Check that directory does not exist before creating it}'' require performing multi-step reasoning on code.

We suggest overcoming these challenges by introducing a modular workflow based on the semantic structure of the query. 
Our approach is based on the intuition of how an engineer would approach a SCS task. For example, in performing search for code that navigates folders in Python they would first only pay attention to code that has cues about operating with paths, directories or folders. Afterwards, they would seek indications of iterating through some of the found objects or other entities in the code related to them. In other words, they would perform multiple steps of different nature - i.e. finding indications of specific types of data entities, or specific operations. Figure~\ref{fig:motivating} illustrates which parts of the code would be important to indicate that they have found the desired code snippet at each step. We attempt to imitate this process in this work. To formalize the decomposition of the query into such steps, we take inspiration from the idea that code is comprised of data, or entities, and transformations, or actions, over data. Thus, a SCS query is also likely to describe the code in terms of data entities and actions.

\begin{figure}[t]
\vspace{-0.8cm}
\centering
\includegraphics[width=0.99\textwidth]{images/overview.pdf}
\vspace{-0.1cm}
\caption{\small \textbf{Overview of the NS3 approach}. We illustrate the pipeline of processing for an example query ``\textit{Load all tables from dataset}''. Parsed query is used for deciding the positions of entity discovery and action modules in the neural module network layout. Each entity discovery module receives a noun/noun phrase as input, and outputs relatedness scores for code tokens, which are passed as input to an action module. Action module gets scores for all its children in the parse-tree, except one, which is masked, and the goal is predicting, cloze-style, what are the relatedness scores for the missing argument.}
\vspace{-0.2cm}
\label{fig:overview}
\end{figure}

We break down the task of matching the query into smaller tasks of matching individual data entities and actions. In particular, we aim to identify parts of the code that indicate the presence of the corresponding data or action. We tackle each part with a distinct type of network -- a neural module. 
Using the semantic parse of the query, we construct the layout of how modules' outputs should be linked according to the relationships between data entities and actions. Correspondingly, this layout specifies how the modules should be combined into a single neural module network (NMN)~\citep{DBLP:conf/cvpr/AndreasRDK16}. NMN is ``evaluated'' on the candidate code snippet to decide whether the query and code are related. This pipeline is outlined in Figure~\ref{fig:overview}, and described in further detail in Section~\ref{sec:model}.


This approach has the following advantages. First, semantic parse captures the compositionality of a query. Second, it mitigates the challenges of encoding a long text by removing the need for it, as each of the modules handle only a small portion of the query. A smaller portion of the query is less likely to require compositional reasoning from an individual neural module. Finally, applying the neural modules in a succession can potentially mimic staged reasoning necessary for SCS. 

We evaluate our proposed NS$^3$ model on two SCS datasets - CodeSearchNet (CSN) \citep{DBLP:journals/corr/abs-1909-09436} and CoSQA/WebQueryTest \citep{DBLP:conf/acl/HuangTSGXJZD21}. Additionally, we experiment with a limited training set size of CSN of 10K and 5K examples. We find that NS$^3$ provides large improvements upon baselines in all cases. To measure how well the model captures compositional properties of the queries, we experiment with performing adversarial modifications to the query and showing that compared to the baselines, NS$^3$ is more likely to correctly recognize modified queries. We also research the effect on queries that require multi-staged reasoning over code by breaking down the performance by query's compositional complexity, and demonstrate that our proposed method handles this task better than the baselines.

Our main contributions are:  
(i) We propose looking at SCS as a compositional task that requires multi-step reasoning.
(ii) We present an implementation of the aforementioned paradigm based on NMNs. 
(iii) We demonstrate that our proposed model provides a large improvement on a number of well-established baseline models.
(iv) We perform additional studies to evaluate the capacity of our model to handle compositional queries. 

\section*{References}


\section*{Checklist}




\begin{enumerate}


\item For all authors...
\begin{enumerate}
  \item Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
    \answerYes{}
  \item Did you describe the limitations of your work?
    \answerYes{See the Discussion paragraph under Section~\ref{sec:analysis-and-case-study}}
  \item Did you discuss any potential negative societal impacts of your work?
    \answerNA{}
  \item Have you read the ethics review guidelines and ensured that your paper conforms to them?
    \answerYes{}
\end{enumerate}


\item If you are including theoretical results...
\begin{enumerate}
  \item Did you state the full set of assumptions of all theoretical results?
    \answerNA{}
        \item Did you include complete proofs of all theoretical results?
    \answerNA{}
\end{enumerate}


\item If you ran experiments...
\begin{enumerate}
  \item Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
    \answerYes{Code and data in supplemental material}
  \item Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
    \answerYes{See paragraph titled Hyperparameters in Section~\ref{sec:experiment-setting}}
        \item Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
    \answerNo{}
        \item Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
    \answerYes{}
\end{enumerate}


\item If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
\begin{enumerate}
  \item If your work uses existing assets, did you cite the creators?
    \answerYes{}
  \item Did you mention the license of the assets?
    \answerNo{}
  \item Did you include any new assets either in the supplemental material or as a URL?
    \answerNo{}
  \item Did you discuss whether and how consent was obtained from people whose data you're using/curating?
    \answerNA{}
  \item Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
    \answerNA{}
\end{enumerate}


\item If you used crowdsourcing or conducted research with human subjects...
\begin{enumerate}
  \item Did you include the full text of instructions given to participants and screenshots, if applicable?
    \answerNA{}
  \item Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
    \answerNA{}
  \item Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
    \answerNA{}
\end{enumerate}


\end{enumerate}





\section{Neural Modular Code Search}
\label{sec:model}
\vspace{-0.2cm}

We propose to supplement the query with a loose structure resembling its semantic parse, as illustrated in Figure~\ref{fig:overview}. We follow the parse structure to break down the query into smaller, semantically coherent parts, so that each corresponds to an individual execution step.
The steps are taken in succession by a neural module network composed from a layout that is determined from the semantic parse of the query (Sec.~\ref{sec:layout}).
The neural module network is composed by stacking ``modules'', or jointly trained networks, of distinct types, each carrying out a different functionality.

\vspace{-0.2cm}
\paragraph{Method Overview}
In this work, we define two types of neural modules - \textit{entity discovery} module (denoted by $E$; Sec.~\ref{sec:entity-module}) and \textit{action} module (denoted by $A$; Sec~\ref{sec:action-module}). The entity discovery module estimates semantic relatedness of each code token $c^j_i$ in the code snippet \begin{small}$\mathbf{c}^j=[c^j_1, \ldots, c^j_N]$\end{small} to an entity mentioned in the query -- e.g. ``\textit{all tables}'' or ``\textit{dataset}'' as in  Figure~\ref{fig:overview}.
The action module estimates the likelihood of each code token to be related to an (unseen) entity affected by the action in the query e.g. ``\textit{dataset}'' and ``\textit{load from}'' correspondingly, conditioned on the rest of the input (seen), e.g. ``\textit{all tables}''. The similarity of the predictions of the entity discovery and action modules measures how well the code matches that part of the query.  
The modules are nested - the action modules are taking as input part of the output of another module - and the order of nesting is decided by the semantic parse layout. In the rest of the paper we refer to the inputs of a module as its \textit{arguments}. 

Every input instance fed to the model is a 3-tuple $(\mathbf{q}^i, s_{q^i}, \mathbf{c}^j)$ consisting of a natural language query $\mathbf{q}^i$, the query's semantic parse $s_{q^i}$, a candidate code (sequence) $\mathbf{c}^j$. The goal is producing a binary label $\hat{r}^{ij}=1$ if the code is a match for the query, and 0 otherwise.
The layout of the neural module network, denoted by $L(s_{q^i})$, is created from the semantic structure of the query $s_{q^i}$. 
During inference, given $(\mathbf{q}^i, s_{q^i}, \mathbf{c}^j)$ as input the model instantiates a network based on the layout, passes $\mathbf{q}^i$, $\mathbf{c}^j$ and $s_{q^i}$ as inputs, and obtains the model prediction $\hat{r}^{ij}$. This pipeline is illustrated in Figure~\ref{fig:overview}, and details about creating the layout of the neural module network are presented in Section~\ref{sec:layout}.

During training, we first perform noisy supervision pretraining for both modules. Next, we perform end-to-end training, where in addition to the query, its parse, and a code snippet, the model is also provided a gold output label $r(\mathbf{q}^i, \mathbf{c}^j)=1$ if the code is a match for the query, and $r(\mathbf{q}^i, \mathbf{c}^j)=0$ otherwise. These labels provide signal for joint fine-tuning of both modules (Section~\ref{sec:training}).


\vspace{-0.1cm}
\subsection{Module Network Layout}\label{sec:layout}
\vspace{-0.1cm}
Here we present our definition of the structural representation $s_{q^i}$ for a query $\mathbf{q}^i$, and introduce how this structural representation is used for dynamically constructing the neural module network, i.e. building its layout $L(s_{q^i})$.

\vspace{-0.2cm}
\paragraph{Query Parsing}
To infer the representation $s_{q^i}$, we pair the query (\textit{e.g.}, ``\textit{Load all tables from dataset}'', as in Figure~\ref{fig:overview}), with a simple semantic parse that looks similar to: \begin{small}\texttt{DO WHAT [ (to/from/in/\ldots) WHAT, WHEN, WHERE, HOW, etc]}\end{small}. 
Following this semantic parse, we break down the query into shorter semantic phrases using the roles of different parts of speech. Nouns and noun phrases correspond to data entities in code, and verbs describe actions or transformations performed on the data entities. Thus, data and transformations are separated and handled by separate neural modules -- an \textit{entity discovery module} $E$ and an \textit{action module} $A$. 
We use a Combinatory Categorial Grammar-based (CCG) semantic parser~\citep{DBLP:journals/corr/abs-1207-1420,DBLP:conf/emnlp/ArtziLZ15} to infer the  semantic parse  $s_{q^i}$ for the natural language query $\mathbf{q}^i$. Parsing is described in further detail in Section~\ref{sec:parser} and Appendix~\ref{sec:app-parsing}.



\vspace{-0.2cm}
\paragraph{Specifying Network Layout}
In the layout \begin{small}$L(s_{q^i})$\end{small}, every noun phrase (e.g., ``\textit{dataset}" in Figure~\ref{fig:overview}) will be passed through the entity discovery module $E$. Module $E$ then produces a probability score $e_k$ for every token $c^j_k$ in the code snippet $\mathbf{c}^j$ to indicate its semantic relatedness to the noun phrase:
\begin{small}
$E(\text{``\textit{dataset}''}, \mathbf{c}^j) = [e_1, e_2, \ldots, e_N].$
\end{small}
Each verb in $s_{q^i}$ (e.g., ``\textit{load}'' in Figure~\ref{fig:overview}) will be passed through an action module:
\begin{small}$A(\text{``\textit{load}''}, \mathbf{p}^i, \mathbf{c}^j) = [a_1, a_2, \ldots, a_N]$\end{small}.
Here, $\mathbf{p}^i$ is the span of arguments to the verb (action) in query $\mathbf{q}^i$, consisting of children of the verb in the parse $s_{q^i}$ (e.g. subject and object arguments to the predicate \textit{``load''}); $a_1, \ldots, a_N$ are estimates of the token scores $e_1, \ldots, e_N$ for an entity from $\mathbf{p}^i$.
The top-level of the semantic parse is always an action module. Figure~\ref{fig:overview} also illustrates preposition \textsc{FROM} used with ``\textit{dataset}'', handling which is described in Section~\ref{sec:action-module}.

\vspace{-0.1cm}
\subsection{Entity Discovery Module}
\label{sec:entity-module}
\vspace{-0.1cm}


The entity discovery module receives a string that references a data entity. Its goal is to identify tokens in the code that have high relevance to that string. 
The architecture of the module is shown in Figure~\ref{fig:entitydiscovery}. Given an entity string, ``\textit{dataset}'' in the example, and a sequence of code tokens $[c^j_1, \ldots, c^j_N]$, entity module first obtains contextual code token representation using RoBERTa model that is initialized from CodeBERT-base checkpoint. The resulting embedding is passed through a two-layer MLP to obtain a score for every individual code token $c^j_k$ : $0 \leq e_k \leq 1$. Thus, the total output of the module is a vector of scores: $[e_1, e_2, \ldots, e_N]$. 
To prime the entity discovery module for measuring relevancy between code tokens and input, we fine-tune it with noisy supervision, as detailed below. 

\vspace{-0.2cm}
\paragraph{Noisy Supervision}
We create noisy supervision for the entity discovery module by using keyword matching and a Python static code analyzer.
For the keyword matching, if a code token is an exact match for one or more tokens in the input string, its supervision label is set to 1, otherwise it is 0. Same is true if the code token is a substring or a superstring of one or more input string tokens. For some common nouns we include their synonyms (e.g. ``\textit{map}'' for ``\textit{dict}''), the full list of those and further details are presented in Appendix~\ref{sec:app-entity-module}.

\begin{wrapfigure}{r}{0.45\textwidth}
\centering
\vspace{-0.3cm}
\includegraphics[width=0.45\textwidth]{images/entity-module}
\vspace{-0.5cm}
\caption{\small Entity module architecture.}
\label{fig:entitydiscovery}
\vspace{-0.4cm}
\end{wrapfigure}

We used the static code analyzer to extract information about statically known data types. We cross-matched this information with the query to discover whether the query references any datatypes found in the code snippet. If that is the case, the corresponding code tokens are assigned supervision label 1, and all the other tokens are assigned to $0$. 
In the pretraining we learned on equal numbers of $(query, code)$ pairs from the dataset, as well as randomly mismatched pairs of queries and code snippets to avoid creating bias in the entity discovery module. 

\vspace{-0.2cm}
\subsection{Action Module}
 \label{sec:action-module}
\vspace{-0.1cm}

First, we discuss the case where the action module has only entity module inputs.
Figure~\ref{fig:action-estimation} provides a high-level illustration of the action module. 
In the example, for the query ``\emph{Load all tables from dataset}'', the action module receives only part of the full query -- ``\emph{Load all tables from ???}''. 
Action module then outputs token scores for the masked argument -- ``\emph{dataset}''. 
If the code snippet corresponds to the query, then the action module should be able to deduce this missing part from the code and the rest of the query. For consistency, we always mask the last data entity argument. 
We pre-train the action module using the output scores of the entity discovery module as supervision.

\begin{wrapfigure}{r}{0.47\textwidth}
\vspace{-0.5cm}
\centering
\includegraphics[width=0.47\textwidth]{images/action-module}
\vspace{-0.5cm}
\caption{\small Action module architecture.}
\label{fig:action-estimation}
\vspace{-0.3cm}
\end{wrapfigure}
More specifically, after masking one input, we compute joint embeddings for verb and preposition for each data entity argument. This is done with a 2-layer MLP model, as illustrated in the left-most part of Figure~\ref{fig:action-estimation}. 
If a data entity does not have a preposition associated with it, that part of the input is filled with zeros. The joint verb-preposition embedding is stacked with the code token embedding $t_k$ and entity discovery module output for that token, this is referenced in the middle part of Figure~\ref{fig:action-estimation}. This vector is passed through a transformer encoder model, followed by a 2-layer MLP and a sigmoid layer to output token score $a_k$, illustrated in the right-most part of the Figure~\ref{fig:action-estimation}.

\vspace{-0.1cm}
\subsection{Model Prediction}
\vspace{-0.2cm}

The final score $\hat{r}^{ij}=f(\mathbf{q}^i, \mathbf{c}^j)$ is computed based on the similarity of action and entity discovery module output scores. Formally, for an action module with verb $x$ and parameters $\mathbf{p}^x = [p^x_1, \ldots, p^x_k]$, the final model prediction is the dot product of respective outputs of action and entity discovery modules: $\hat{r}^{ij} = A(x, p^x_1, \ldots, p^x_{k-1}) \cdot E(p^x_k)$.
Since the action module estimates token scores for the entity affected by the verb, if its prediction is far from the truth - then either the action is not found in the code, or it is not fully corresponding to the query, for example, in the code snippet tables are loaded from web, instead of a dataset.
We normalize this score to make it a probability. If this is the only action in the query, this probability score will be the output of the entire model for ($\mathbf{q}^i, \mathbf{c}^j$)  pair: $\hat{r}^{ij}$, otherwise $\hat{r}^{ij}$ will be the product of probability scores of all nested actions in the layout.


\vspace{-0.2cm}
\paragraph{Compositional query with nested actions} 
Consider a compositional query ``\textit{Load all tables from dataset using \emph{Lib} library}''. Here action with verb ``\textit{Load from}'' has an additional argument ``\textit{using}'' -- also an action -- with an entity argument ``\textit{\emph{Lib} library}''. In case of nested actions, 
we flatten the layout by taking the conjunction of individual action similarity scores. 
Formally, for two verbs $x$ and $y$ and their corresponding arguments $\mathbf{p}^x = [p^x_1, \ldots, p^x_k]$ and $\mathbf{p}^y = [p^y_1, \ldots, p^y_l]$ in a layout that looks like: $A(x, \mathbf{p}^x, A(y, \mathbf{p}^y))$, 
the output of the model is the conjunction of similarity scores computed for individual action modules:
$sim(A(x, p^x_1, \ldots, p^x_{k-1}), E(p^x_k)) \cdot sim(A(y, p^y_1, \ldots, p^y_{l-1}), E(p^y_l))$. This process is repeated until all remaining $\mathbf{p}^x$ and $\mathbf{p}^y$ are data entities. 
This design ensures that code snippet is ranked highly if both actions are ranked highly, we leave explorations of alternative handling approaches for nested actions to future work.

\vspace{-0.2cm}
\subsection{Module Pretraining and Joint Fine-tuning}
\label{sec:training}
\vspace{-0.1cm}
We train our model through supervised pre-training, as is discussed in Sections~\ref{sec:entity-module} and~\ref{sec:action-module}, followed by end-to-end training. End-to-end training objective is binary classification - given a pair of query $\mathbf{q}^i$ and code $\mathbf{c}^j$, the model predicts probability $\hat{r}^{ij}$ that they are related. In the end-to-end training, we use positive examples taken directly from the dataset - ($\mathbf{q}^i$, $\mathbf{c}^i$), as well as negative examples composed through the combination of randomly mismatched queries and code snippets. The goal of end-to-end training is fine-tuning parameters of entity discovery and action modules, including the weights of the RoBERTA models used for code token representation. 

Batching is hard to achieve for our model, so for the interest of time efficiency we do not perform inference on all distractor code snippets in the code dataset. Instead, for a given query we re-rank top-K highest ranked code snippets as outputted by some baseline model, in our evaluations we used CodeBERT. Essentially, we use our model in a re-ranking setup, this is common in information retrieval and is known as L2 ranking. 
We interpret the probabilities outputted by the model as ranking scores. More details about this procedure are provided in Section~\ref{sec:eval-metrics}.


\section{Background}
\vspace{-0.2cm}

\subsection{Semantic Code Search}
\vspace{-0.2cm}
Semantic code search (SCS) is the process of retrieving a relevant code snippet based on a textual description of its functionality, also referred to as \textit{query}.
Let $\mathcal{C}$ be a database of code snippets $\mathbf{c}^i$. For each $\mathbf{c}^i \in \mathcal{C}$, there is a textual description of its functionality $\mathbf{q}^i$. In the example in Figure~\ref{fig:overview}, the query $\mathbf{q}^i$ is ``\textit{Load all tables from dataset}''.
Let $r$ be an indicator function such that $r(\mathbf{q}^i, \mathbf{c}^j) = 1$ if $i=j$; and 0 otherwise.
Given some query $\mathbf{q}$ the goal of SCS is to find $\mathbf{c}^*$ such that $r(\mathbf{q}, \mathbf{c}^*) = 1$. We assume that for each $\mathbf{q}^*$ there is exactly one such $\mathbf{c}^*$.\footnote{This is not the case in CoSQA dataset. For the sake of consistency, we perform the evaluation repeatedly, leaving only one correct code snippet among the candidates at a time, while removing the others.}
Here we look to construct a model which takes as input a pair of query and a candidate code snippet: $(\mathbf{q}^i, \mathbf{c}^j)$ and assign the pair a probability $\hat{r}^{ij}$ for being a correct match. Following the common practice in information retrieval, we evaluate the performance of the model based on how high the correct answer $\mathbf{c}^*$ is ranked among a number of incorrect, or distractor instances $\{\mathbf{c}\}$. This set of distractor instances can be the entire codebase $\mathcal{C}$, or a subset of the codebase obtained through heuristic filtering, or another ranking method. 

\vspace{-0.1cm}
\subsection{Neural Models for Semantic Code Search}
\vspace{-0.1cm}

Past works handling programs and code have focused on enriching their models with incorporating more semantic and syntactic information from code~\citep{DBLP:conf/acl/AhmadCRC20,DBLP:conf/acl/ChoiBNL21,DBLP:conf/emnlp/Shi0D0HZS21,DBLP:conf/iclr/ZugnerKCLG21}. 
Some prior works have cast the SCS as a sequence classification task, where the code is represented as a textual sequence and input pair $(\mathbf{q}^i, \mathbf{c}^j)$ is concatenated with a special separator symbol into a single sequence, and the output is the score $\hat{r}^{ij}$: $\hat{r}^{ij} = f(\mathbf{q}^i, \mathbf{c}^j)$. The function $f$ performing the classification can be any sequence classification model, e.g. BERT~\citep{DBLP:conf/naacl/DevlinCLT19}. 

Alternatively, one can define separate networks for independently representing the query ($f$), the code ($g$) and measuring the similarity between them: $\hat{r}^{ij} = sim(f(\mathbf{q}^i), g(\mathbf{c}^j))$. This allows one to design the code encoding network $g$ with additional program-specific information, such as abstract syntax trees~\citep{DBLP:journals/pacmpl/AlonZLY19, DBLP:conf/icse/ZhangWZ0WL19} or control flow graphs~\citep{DBLP:journals/nn/GuLGWZXL21, DBLP:conf/sigsoft/ZhaoH18}. Separating two modalities of natural language and code also allows further enrichment of code representation by adding contrastive learning objectives~\citep{DBLP:conf/emnlp/0001JZA0S21, DBLP:conf/sigir/BuiYJ21}. In these approaches, the original code snippet $\mathbf{c}$ is automatically modified with semantic-preserving transformations, such as variable renaming, to introduce versions of the code snippet - $\mathbf{c}^\prime$ with the exact same functionality. Code encoder $g$ is then trained with an appropriate contrastive loss, such as Noise Contrastive Estimation (NCE)~\citep{DBLP:journals/jmlr/GutmannH10}, or InfoNCE~\citep{DBLP:journals/corr/abs-1807-03748}. 



\vspace{-0.0cm}
\textbf{Limitations}
However, there is also merit in reviewing how we represent and use the textual \textit{query} to help guide the SCS process.
Firstly, existing work derive a single embedding for the entire query. The problem of representing longer text inputs is being actively researched in the community and remains an open problem~\citep{DBLP:journals/corr/abs-2004-05150,DBLP:conf/rep4nlp/BhathenaWD20, DBLP:journals/corr/abs-1904-00157}. This means that specific details or nested subqueries of the query may be omitted or not represented properly - getting lost in the embedding. Secondly, prior approaches make the decision after a single pass over the code snippet. 
This ignores cases where reasoning about a query requires multiple steps and thus - multiple look-ups over the code, as is for example in cases with nested subqueries.
Our proposed approach - $NS^3$ - attempts to address these issues by breaking down the query into smaller phrases based on its semantic parse and locating each of them in the code snippet. This should allow us to match compositional and longer queries to code more precisely.
\section{Related work}
\vspace{-0.2cm}
Different deep learning models have proved quite efficient when applying to programming languages and code. Prior works have studied and reviewed the uses of deep learning for code analysis in general and code search in particular~\citep{DBLP:journals/corr/abs-2101-11149, DBLP:conf/nips/LuGRHSBCDJ21}.

A number of approaches to deep code search is based on creating a relevance-predicting model between text and code. 
\citep{DBLP:conf/icse/GuZ018} propose using RNNs for embedding both code and text to the same latent space. 
On the other hand, \citep{DBLP:journals/tkdd/LingWWPMXLWJ21} capitalizes the inherent graph-like structure of programs to formulate code search as graph matching.
A few works propose enriching the models handling code embedding by adding additional code analysis information, such as semantic and dependency parses~\citep{DBLP:conf/cikm/DuSWSHZ21, DBLP:conf/msr/AkbarK19}, variable renaming and statement permutation~
\citep{DBLP:journals/corr/abs-2012-01028}, as well as structures such as abstract syntax tree of the program~\citep{DBLP:conf/acl/HaldarWXH20, DBLP:conf/kbse/WanSSXZ0Y19}.
A few other approaches have dual formulations of code retrieval and code summarization~\citep{DBLP:conf/kbse/ChenZ18,  DBLP:conf/www/YaoPS19, DBLP:conf/www/YeXZHWZ20,DBLP:conf/sigir/BuiYJ21}
In their work \citep{DBLP:journals/corr/abs-2008-12193} propose considering the code search scenario where short annotative descriptions of code snippets are provided. Appendix~\ref{sec:app-related-work} discusses more related work.

\section{Experiments}
\vspace{-0.2cm}

\subsection{Experiment Setting}
\label{sec:experiment-setting}
\vspace{-0.1cm}


\paragraph{Dataset}
We conduct experiments on two datasets: Python portion of the \textbf{CodeSearchNet} (CSN) \citep{DBLP:journals/corr/abs-1909-09436}, and \textbf{CoSQA} \citep{DBLP:conf/acl/HuangTSGXJZD21}. We parse all queries with the CCG parser, Section~\ref{sec:parser}, excluding unparsable examples from further experiments. This leaves us with approximately 40\% of the CSN dataset and 70\% of the CoSQA dataset, the exact data statistics are available in Appendix~\ref{sec:app-experiment-settings} in Table~\ref{tab:data-size}. We believe, that the difference in success rate of the parser between the two datasets can be attributed to the fact that CSN dataset, unlike CoSQA, does not contain real code search queries, but rather consists of docstrings, which are used as approximate queries. More details and examples can be found in Appendix~\ref{sec:app-docstrings}.
For our baselines, we use the parsed portion of the dataset for fine-tuning to make the comparison fair. In addition, we also experiment with fine-tuning all models on an even smaller subset of CodeSearchNet dataset, using only 5K and 10K examples for fine-tuning. The goal is testing whether modular design makes NS3 more sample-efficient.

All experiment and ablation results discussed in Sections~\ref{sec:results},\ref{sec:ablation-studies} and \ref{sec:analysis-and-case-study} are obtained on the test set of CSN for models trained on CSN training data, or WebQueryTest \citep{DBLP:conf/nips/LuGRHSBCDJ21} --  a small natural language web query dataset of document-code pairs -- for models trained on CoSQA dataset.

\vspace{-0.2cm}
\paragraph{Evaluation and Metrics}\label{sec:eval-metrics}
We follow CodeSearchNet's original approach for evaluation for a test instance $(q, c)$, comparing the output against outputs over a fixed set of 999 distractor code snippets.
We use two evaluation metrics: Mean Reciprocal Rank (MRR) and Precision@K (P@K) for K=1, 3, and 5, see Appendix~\ref{sec:app-metrics} for definitions and further details.

Following a common approach in information retrieval, we perform two-step evaluation. In the first step, we obtain CodeBERT's output against 999 distractors. In the second step, we use $NS^3$ to re-rank the top 10 predictions of CodeBERT. This way the evaluation is much faster, since unlike our modular approach, CodeBERT can be fed examples in batches. And as we will see from the results, we see improvement in final performance in all scenarios. 


\vspace{-0.2cm}
\paragraph{Compared Methods}
We compare $\textsc{NS}^3$ with various state-of-the-art methods, including some traditional approaches for document retrieval and pretrained large NLP language models. (1) \textbf{BM25} is a ranking method to estimate the relevance of documents to a given query. 
(2) \textbf{RoBERTa (code)} is a variant of RoBERTa \citep{DBLP:journals/corr/abs-1907-11692} pretrained on the CodeSearchNet corpus.
(3) \textbf{CuBERT} \citep{DBLP:conf/icml/KanadeMBS20} is a BERT Large model pretrained on 7.4M Python files from GitHub. 
(4) \textbf{CodeBERT} is an encoder-only Transformer model trained on unlabeled source code via masked language modeling (MLM) and replaced token detection objectives. 
(5) \textbf{GraphCodeBERT} is a pretrained Transformer model using MLM, data flow edge prediction, and variable alignment between code and the data flow. 
(6) \textbf{GraphCodeBERT*} is a re-ranking baseline. We used the same setup as for NS3, but used GraphCodeBERT to re-rank the top-10 predictions of the CodeBERT model. 

\vspace{-0.2cm}
\paragraph{Query Parser}\label{sec:parser}
We started by building a vocabulary of predicates for common action verbs and entity nouns, such as ``\textit{convert}'', ``\textit{find}'', ``\textit{dict}'', ``\textit{map}'', etc.
For those we constructed the lexicon (rules) of the parser. We have also included ``catch-all'' rules, for parsing sentences with less-common words. To increase the ratio of the parsed data, we preprocessed the queries by removing preceding question words, punctuation marks, etc. Full implementation of our parser including the entire lexicon and vocabulary can be found at \url{https://anonymous.4open.science/r/ccg_parser-4BC6}. Details are available in Appendix~\ref{sec:app-parsing}.

\vspace{-0.2cm}
\paragraph{Pretrained Models}
Action and entity discovery modules each embed code tokens with a RoBERTa model, that has been initialized from a checkpoint of pretrained CodeBERT model \footnote{\url{https://huggingface.co/microsoft/codebert-base}}. We fine-tune these models during the pretraining phases, as well as during final end-to-end training phase. 


\begin{wraptable}{r}{0.43\textwidth}
\vspace{-0.4cm}
    \hspace{-0.1cm}
    \scalebox{0.76}{
        \begin{tabular}{lllll} \\
        \toprule
        \multirow{2}{*}{Method} & \multicolumn{4}{c}{CoSQA} \\ 
        & MRR & P@1 & P@3 & P@5 \\
        \midrule
        BM25 & 0.103 & 0.05 & 0.119 & 0.142\\
       
        RoBERTa (code) & 0.279  & 0.159 & 0.343 & 0.434 \\
        CuBERT & 0.127 & 0.067 & 0.136 & 0.187 \\
        CodeBERT & 0.345 & 0.175 &  0.42 & 0.54 \\
        GraphCodeBERT & 0.435 & 0.257 & 0.538 & 0.628\\
        GraphCodeBERT* & 0.462 & 0.314 & 0.547 & 0.632\\
        \midrule
        NS$^3$  & \textbf{0.551} & \textbf{0.445} & \textbf{0.619} & \textbf{0.668}\\
        \midrule
        Upper-bound & 0.736 & 0.724 & 0.724 & 0.724\\
        \bottomrule
        \end{tabular}
    }
    \caption{\small Mean Reciprocal Rank(MRR) and Precision@1/@3/@5 (higher is better) for different methods trained on CoSQA dataset.}
     \label{tab:cosqa}
    \vspace{-0.2cm}
\end{wraptable}

\begin{table*}[t]
\vspace{-0.7cm}
    \centering
    \scalebox{0.78}{
        \begin{tabular}{lllllllllllll} \\
        \toprule
            \multirow{2}{*}{Method} & \multicolumn{4}{c}{CSN} & \multicolumn{4}{c}{CSN-10K} & \multicolumn{4}{c}{CSN-5K} \\ 
             & MRR & P@1 & P@3 & P@5 & MRR & P@1 & P@3 & P@5 & MRR & P@1 & P@3 & P@5\\ 
            \midrule
            BM25 & 0.209 & 0.144 & 0.230 & 0.273 & 0.209 & 0.144 & 0.230 & 0.273 & 0.209 & 0.144 & 0.230 & 0.273 \\
           
            RoBERTa (code) & 0.842 & 0.768 & 0.905 & 0.933 & 0.461 & 0.296 & 0.545 & 0.664 & 0.290 & 0.146 & 0.324 & 0.438\\
            CuBERT & 0.225 & 0.168  & 0.253 & 0.294 & 0.144 & 0.081 & 0.166 & 0.214 & 0.081 & 0.030 & 0.078 & 0.118 \\
            CodeBERT & 0.873 & 0.803 & 0.939 & 0.958 & 0.69 & 0.55 & 0.799 & 0.873 & 0.680 & 0.535 & 0.794 & 0.870\\
            GraphCodeBERT & 0.812 & 0.725 & 0.880 & 0.919 & 0.786 & 0.684 & 0.859 & 0.901 & 0.773 & 0.677 & 0.852 & 0.892\\
            GraphCodeBERT* & 0.883 & 0.820 & 0.941 & 0.962 & 0.780 & 0.683 & 0.858 & 0.904 & 0.765 & 0.662 & 0.846 & 0.894\\
            \midrule
            NS$^3$ & \textbf{0.924} & \textbf{0.884} & \textbf{0.959} & \textbf{0.969} & \textbf{0.826} & \textbf{0.753} & \textbf{0.886} & \textbf{0.908} & \textbf{0.823}    & \textbf{0.751} & \textbf{0.881} & \textbf{0.913} \\
            \midrule
            Upper-bound & 0.979 &  &  &  & 0.939 &  &  &  & 0.936 &   &   &  \\
            \bottomrule
        \end{tabular}
    }
    \caption{Mean Reciprocal Rank (MRR) and Precision@1/@3/@5 (higher is better) for methods trained on different subsets from CodeSearchNet dataset.}
    \label{tab:results-csn}
    \vspace{-0.1cm}
\end{table*}

\vspace{-0.2cm}
\paragraph{Hyperparameters}
The MLPs in entity discovery and action modules have 2 layers with input dimension of 768. We use dropout in these networks with rate $0.1$. The learning rate for pretraining and end-to-end training phases was chosen from the range of 1e-6 to 6e-5. We use early stopping with evaluation on unseen validation set for model selection during action module pretraining and end-to-end training. For entity discovery model selection we performed manual inspection of produced scores on unseen examples. For fine-tuning the CuBERT, CodeBERT and GraphCodeBERT baselines we use the hyperparameters reported in their original papers. For RoBERTa (code), we perform the search for learning rate during fine-tuning stage in the same interval as for our model. For model selection on baselines we also use early stopping.


\vspace{-0.1cm}
\subsection{Results}\label{sec:results}
\vspace{-0.2cm}

\paragraph{Performance Comparison}
Tables~\ref{tab:results-csn}~and~\ref{tab:cosqa} present the performance evaluated on testing portion of CodeSearchNet dataset, and WebQueryTest dataset correspondingly. As it can be seen, our proposed model outperforms the baselines. 

Our evaluation strategy improves performance only if the correct code snippet was ranked among the top-10 results returned by the CodeBERT model, so rows labelled ``Upper-bound'' report best possible performance with this evaluation strategy. 
 
\begin{figure*}[t]
\begin{subfigure}[h]{0.45\textwidth}
\centering
\includegraphics[width=\textwidth]{images/barchart-table6}
\caption{}
\label{fig:max-depth}
\end{subfigure}
\begin{subfigure}[h]{0.52\textwidth}
\centering
\includegraphics[width=\textwidth]{images/barchart-table7}
\caption{}
\end{subfigure}
\vspace{-0.2cm}
\caption{\small We report Precision@1 scores. (a) Performance of our proposed method and baselines broken down by average number of arguments per action in a single query. (b) Performance of our proposed method and baselines broken down by number of arguments in queries with a single action.}
\label{fig:arguments-number}
\vspace{-0.1cm}
\end{figure*}

\paragraph{Query Complexity vs. Performance}
Here we present the breakdown of the performance for our method vs baselines, using two proxies for the complexity and compositionality of the query. The first one is the maximum depth of the query. We define the maximum depth as the maximum number of nested action modules in the query. The results for this experiment are presented in Figure~\ref{fig:max-depth}. As we can see, $NS^3$ improves over the baseline in all scenarios. It is interesting to note, that while CodeBERT achieves the best performance on queries with depth 3+, our model's performance peaks at depth = 1. We hypothesize that this can be related to the automated parsing procedure, as parsing errors are more likely to be propagated in deeper queries. Further studies with carefully curated manual parses are necessary to better understand this phenomenon. 

\begin{figure*}[t]
\centering
\begin{subfigure}[h]{0.31\textwidth}
\centering
\includegraphics[width=\textwidth]{images/training-ablations.pdf}
\caption{Effect of training}
\label{fig:training-ablations}
\end{subfigure}
\begin{subfigure}[h]{0.33\textwidth}
\centering
\includegraphics[width=\textwidth]{images/norm-ablations.pdf}
\caption{Score normalization}
\label{fig:norm-ablations}
\end{subfigure}
\begin{subfigure}[h]{0.32\textwidth}
\centering
\includegraphics[width=\textwidth]{images/similarity-ablations}
\caption{Similarity measure}
\label{fig:similarity-ablations}
\end{subfigure}
\vspace{-0.1cm}
\caption{\small Performance of $NS^3$ on the test portion of CSN dataset with different ablation variants. (a) Skipping one, or both pretraining procedures, and only training end-to-end. (b) Using no normalization on output scores (\textbf{None}), action-only or entity discovery-only, and both. (c) Performance with different options for computing action and entity discovery output similarities.}
\label{fig:ablations}
\vspace{-0.4cm}
\end{figure*}

Another proxy for the query complexity we consider, is the number of data arguments to a single action module. While the previous scenario is breaking down the performance by the depth of the query, here we consider its ``width''. We measure the average number of entity arguments per action module in the query. In the parsed portion of our dataset we have queries that range from 1 to 3 textual arguments per action verb. The results for this evaluation are presented in Figure~\ref{fig:arguments-number}. As it can be seen, there is no significant difference in performances between the two groups of queries in either CodeBERT or our proposed method - $NS^3$.

\vspace{-0.1cm}
\subsection{Ablation Studies}\label{sec:ablation-studies}
\vspace{-0.2cm}
\paragraph{Effect of Pretraining}
In an attempt to better understand the individual effect of the two modules as well as the roles of their pretraining and training procedures, we performed two additional ablation studies. In the first one, we compare the final performance of the original model with two versions where we skipped part of the pretraining. The model noted as ($NS^3-AP$) was trained with pretrained entity discovery module, but no pretraining was done for action module, instead we proceeded to the end-to-end training directly. For the model called $NS^3-(AP\&EP)$, we skipped both pretrainings of the entity and action modules, and just performed end-to-end training. Figure~\ref{fig:training-ablations} demonstrates that combined pretraining is important for the final performance. 
Additionally, we wanted to measure how effective the setup was without end-to-end training. The results are reported in Figure~\ref{fig:training-ablations} under the name $NS^3-E2E$. There is a huge performance dip in this scenario, and while the performance is better than random, it is obvious that end-to-end training is crucial for $NS^3$.

\vspace{-0.3cm}
\subparagraph{Score Normalization}
We wanted to determine the importance of output normalization for the modules to a proper probability distribution. In Figure~\ref{fig:norm-ablations} we demonstrate the performance achieved using no normalization at all, normalizing either action or entity discovery module, or normalizing both. In all cases we used L1 normalization, since our output scores are non-negative. The version that is not normalized at all performs the worst on both datasets. The performances of the other three versions are close on both datasets.

\vspace{-0.3cm}
\subparagraph{Similarity Metric}
Additionally, we experimented with replacing the dot product similarity with a different similarity metric. In particular, in Figure~\ref{fig:similarity-ablations} we compare the performance achieved using dot product similarity, L2 distance, and weighted cosine similarity. The difference in performance among different versions is marginal.

\vspace{-0.1cm}
\subsection{Analysis and Case Study}\label{sec:analysis-and-case-study}
\vspace{-0.2cm}

Appendix~\ref{sec:app-additional-analysis} contains additional studies on model generalization, such as handling completely unseen actions and entities, as well as the impact of the frequency of observing an action or entity during training has on model performance. 


\vspace{-0.4cm} 
\subparagraph{Case Study}
Finally, we demonstrate some examples of the scores produced by our modules at different stages of training. Figure~\ref{fig:case-study1} shows module score outputs for two different queries and with their corresponding code snippets. The first column shows the output of the entity discovery module after pretraining, while the second and third columns demonstrate the outputs of entity discovery and action modules after the end-to-end training. We can see that in the first column the model identifies syntactic matches, such as ``folder'' and a list comprehension, which ``elements'' could be related too. After fine-tuning we can see there is a wider range of both syntactic and some semantic matches present, e.g. ``dirlist'' and ``filelist'' are correctly identified as related to ``folders''.

\vspace{-0.4cm}
\paragraph{Perturbed Query Evaluation}
\label{sec:perturbed-query}
In this section we study how sensitive the models are to small changes in the query $\mathbf{q}^i$, so that it no longer correctly describes its corresponding code snippet $\mathbf{c}^i$. 
Our expectation is that evaluating a sensitive model on $\mathbf{c}^i$ will rate the original query higher than the perturbed one. Whereas a model that tends to over-generalize and ignore details of the query will likely rate the perturbed query similar to the original. We start from 100 different pairs $(\mathbf{q}^i, \mathbf{c}^i)$, that both our model and CodeBERT predict correctly. 
\begin{wrapfigure}{r}{0.42\textwidth}
\vspace{-0.4cm}
\centering
\includegraphics[width=0.42\textwidth]{images/barchart-table5}
\caption{\small Ratio of the perturbed query score to the original query score  (lower is better) on CSN dataset.}
\label{fig:perturbed_query_eval}
\vspace{-0.3cm}
\end{wrapfigure}
We limited our study to queries with a single verb and a single data entity argument to that verb. For each pair we generated perturbations of two kinds, with 20 perturbed versions for every query. For the first type of perturbations, we replaced query's data argument with a data argument sampled randomly from another query. For the second type, we replaced the verb argument with another randomly sampled verb.
To account for calibration of the models, we measure the change in performance through ratio of the perturbed query score over original query score (lower is better). 
The results are shown in Figure~\ref{fig:perturbed_query_eval}, labelled ``$V(arg_1) \rightarrow  V(arg_2)$'' and ``$V_1(arg) \rightarrow V_2(arg)$''.


\begin{figure}[t]
\centering
\vspace{-0.7cm}
\includegraphics[width=0.92\textwidth]{images/case-study.png}
\vspace{-0.1cm}
\caption{\small 
Token scores outputted by the modules at different stages of training. Darker highlighting means higher score. The leftmost and middle columns show output scores of the entity discovery module after pretraining, and the end-to-end training correspondingly. The rightmost column shows the scores of the action module after the end-to-end training. }
\label{fig:case-study1}
\vspace{-0.2cm}
\end{figure}

\vspace{-0.3cm}
\paragraph{Discussion}
One of the main requirements for the application of our proposed method is being able to construct a semantic parse of the retrieval query. In general, it is reasonable to expect the users of the SCS to be able to come up with a formal representation of the query, e.g. by representing it in a form similar to SQL or CodeQL. However, due to the lack of such data for training and testing purposes, we implemented our own parser, which understandably does not have perfect performance since we are dealing with open-ended sentences. 
