Code Understanding › Code Search

Code search retrieves relevant code from a large repository given a natural language query. The goal is to find the function, class, or module that implements a given concept, without knowing the exact names or file locations.

Traditional code search is lexical: the search engine finds code files that contain the query terms. This fails when the vocabulary of the query does not match the vocabulary of the code. A query for "how do I read a file line by line" will not match a function named iter_lines unless the function's docstring contains the phrase "line by line."

This vocabulary mismatch is called the semantic gap between natural language and code. It is a fundamental challenge because code authors choose names for reasons of brevity, convention, or domain familiarity, not to match the vocabulary that later users will use to search for the code. A function that implements exponential moving average might be named ema, smooth, running_average, or exp_weighted_mean, and a new user might search for any of these, or for "discount past values" or "weight recent observations more heavily."

Semantic code search addresses this gap by embedding both queries and code into a shared vector space where semantically similar items are close. A natural language query and the code that implements the described behavior should be near each other in this space, even if they share no words. This is a multimodal embedding problem: the encoder must handle both natural language and code as inputs and map them to a common space.

The connection to the pretraining insight from earlier in this chapter is direct: models pretrained on code-documentation pairs have already learned this alignment, because predicting docstrings from code (and code from docstrings) is structurally equivalent to the embedding alignment problem. Fine-tuning with a contrastive objective sharpens this implicit alignment into explicit similarity scores.

The standard architecture for semantic code search is the dual encoder, also called a bi-encoder:

• An NL encoder maps a natural language query $q$ to a vector $\mathbf{q} \in \mathbb{R}^d$.
• A code encoder maps a code snippet $c$ to a vector $\mathbf{c} \in \mathbb{R}^d$.
• Retrieval scores code snippets by $\text{sim}(\mathbf{q}, \mathbf{c})$, typically cosine similarity.

The two encoders may share weights (a single model handles both modalities) or be separate models. CodeBERT and UniXcoder use a shared encoder: the same transformer handles both natural language and code inputs, with the modality distinction emerging from the training distribution. This parameter sharing is efficient and, more importantly, encourages the model to learn representations that are aligned across modalities rather than developing separate representation spaces that must then be bridged.

The two encoders are trained jointly to maximize similarity between matching query-code pairs and minimize similarity between non-matching pairs. The contrastive loss for a batch of $N$ query-code pairs is:

$$\mathcal{L}_{\text{contrastive}} = -\frac{1}{N} \sum_{i=1}^{N} \log \frac{\exp\!\left(\text{sim}(\mathbf{q}_i, \mathbf{c}_i) / \tau\right)}{\sum_{j=1}^{N} \exp\!\left(\text{sim}(\mathbf{q}_i, \mathbf{c}_j) / \tau\right)}$$

where:

• $\text{sim}(\mathbf{q}_i, \mathbf{c}_j) = \frac{\mathbf{q}_i \cdot \mathbf{c}_j}{\|\mathbf{q}_i\| \|\mathbf{c}_j\|}$: cosine similarity between query $i$ and code snippet $j$
• $\tau > 0$: a temperature parameter controlling how sharply the softmax discriminates between close and distant pairs
• The numerator corresponds to the correct pair $(i, i)$; the denominator sums over all code snippets in the batch as negatives

This is the InfoNCE loss, the same objective used in SimCLR for visual representation learning and in CLIP for image-text alignment. The temperature $\tau$ is a critical hyperparameter. When $\tau$ is large, the softmax is flat and the loss signal is weak, because even very dissimilar pairs contribute substantially to the denominator. When $\tau$ is small, the softmax is sharp and the model is penalized heavily for any small misranking. In practice, $\tau$ is learned or set based on a small grid search, with values around 0.07 commonly used for contrastive code search.

The model learns to push the correct query-code pair close together in embedding space while pushing all other code snippets in the batch away from the query. In-batch negatives, where all other code snippets in the same batch serve as negatives for each query, are computationally efficient: each forward pass provides $N-1$ negative examples for each of the $N$ queries in the batch.

At inference time, all code snippets in the repository are encoded once and indexed using an approximate nearest neighbor (ANN) index such as FAISS. A new query is encoded and then searched against the index in sub-linear time.

The quality of the code search model depends heavily on the quality and quantity of query-code pairs used for training. Several sources provide these pairs:

Docstring-code pairs: Function definitions paired with their docstrings are natural query-code pairs. CodeSearchNet assembled 6 million such pairs from GitHub.

Stack Overflow pairs: A natural language question paired with an accepted code answer. Questions on Stack Overflow are often exactly the kind of queries a developer would make to a code search engine.

Synthetic pairs: Large language models can generate queries for code snippets or generate code for queries, expanding the training set beyond what exists in public repositories.

Hard negatives: Training pairs where the negative code snippet is semantically close but not the correct answer. Hard negatives prevent the model from relying on superficial differences and force it to learn fine-grained semantic distinctions. For example, pairing a query "find the minimum element" with both a find_max function (hard negative) and an unrelated string operation (easy negative) forces the encoder to distinguish more precisely between minimum and maximum semantics. Easy negatives alone produce models that can find the right function in a diverse codebase but fail when many functions are semantically similar.

Dual-encoder retrieval is fast but imprecise at fine-grained ranking. The alternative is a cross-encoder, which processes the query and each candidate code snippet together and produces a scalar score. Cross-encoders can compute arbitrary interactions between query tokens and code tokens (because both appear in the same forward pass), making them much more accurate than dual encoders at fine-grained ranking. But they cannot be pre-computed: every query requires a forward pass against every candidate snippet, making them impractical for large codebases when used alone.

The standard solution is two-stage retrieval. The dual encoder retrieves the top-K candidates quickly (typically $K = 100$). The cross-encoder then re-ranks those $K$ candidates, selecting the top few results to return to the developer. This hybrid approach gets the recall benefits of fast approximate retrieval with the precision benefits of expensive joint scoring. Products like GitHub Code Search use architectures of this general form.